diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index 4e72a7dbcf..1596db26b2 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -38,11 +38,11 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v2 + uses: actions/checkout@v3 # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL - uses: github/codeql-action/init@v1 + uses: github/codeql-action/init@v2 with: languages: ${{ matrix.language }} # If you wish to specify custom queries, you can do so here or in a config file. @@ -53,7 +53,7 @@ jobs: # 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@v1 + # uses: github/codeql-action/autobuild@v2 # ℹī¸ Command-line programs to run using the OS shell. # 📚 https://git.io/JvXDl @@ -68,4 +68,4 @@ jobs: BUILD_MONAI=1 ./runtests.sh --build - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@v1 + uses: github/codeql-action/analyze@v2 diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml index 73098cd63d..98c194f474 100644 --- a/.github/workflows/conda.yml +++ b/.github/workflows/conda.yml @@ -30,7 +30,7 @@ jobs: minimum-size: 8 maximum-size: 16 disk-root: "D:" - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - uses: conda-incubator/setup-miniconda@v2 with: auto-update-conda: true diff --git a/.github/workflows/cron-mmar.yml b/.github/workflows/cron-mmar.yml index 5fc3ad2d22..46bf7ff384 100644 --- a/.github/workflows/cron-mmar.yml +++ b/.github/workflows/cron-mmar.yml @@ -16,16 +16,16 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: cache weekly timestamp id: pip-cache run: echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ~/.cache/pip diff --git a/.github/workflows/cron.yml b/.github/workflows/cron.yml index f76fe01699..dca7f39617 100644 --- a/.github/workflows/cron.yml +++ b/.github/workflows/cron.yml @@ -15,24 +15,24 @@ jobs: runs-on: [self-hosted, linux, x64, common] strategy: matrix: - pytorch-version: [1.6.0, 1.7.1, 1.8.1, 1.9.1, latest] + pytorch-version: [1.7.1, 1.8.1, 1.9.1, 1.10.2, latest] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Install the dependencies run: | which python python -m pip install --upgrade pip wheel python -m pip uninstall -y torch torchvision if [ ${{ matrix.pytorch-version }} == "latest" ]; then - python -m pip install torch torchvision - elif [ ${{ matrix.pytorch-version }} == "1.6.0" ]; then - python -m pip install torch==1.6.0 torchvision==0.7.0 + python -m pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113 elif [ ${{ matrix.pytorch-version }} == "1.7.1" ]; then - python -m pip install torch==1.7.1 torchvision==0.8.2 + python -m pip install torch==1.7.1 torchvision==0.8.2 --extra-index-url https://download.pytorch.org/whl/cu113 elif [ ${{ matrix.pytorch-version }} == "1.8.1" ]; then - python -m pip install torch==1.8.1 torchvision==0.9.1 + python -m pip install torch==1.8.1 torchvision==0.9.1 --extra-index-url https://download.pytorch.org/whl/cu113 elif [ ${{ matrix.pytorch-version }} == "1.9.1" ]; then - python -m pip install torch==1.9.1 torchvision==0.10.1 + python -m pip install torch==1.9.1 torchvision==0.10.1 --extra-index-url https://download.pytorch.org/whl/cu113 + elif [ ${{ matrix.pytorch-version }} == "1.10.2" ]; then + python -m pip install torch==1.10.2 torchvision==0.11.3 --extra-index-url https://download.pytorch.org/whl/cu113 fi python -m pip install -r requirements-dev.txt python -m pip list @@ -62,13 +62,13 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' strategy: matrix: - container: ["pytorch:21.02", "pytorch:21.10", "pytorch:22.03"] # 21.02, 21.10 for backward comp. + container: ["pytorch:21.02", "pytorch:21.10", "pytorch:22.06"] # 21.02, 21.10 for backward comp. container: image: nvcr.io/nvidia/${{ matrix.container }}-py3 # testing with the latest pytorch base image options: "--gpus all" runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Install APT dependencies run: | apt-get update @@ -106,13 +106,13 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' strategy: matrix: - container: ["pytorch:21.02", "pytorch:21.10", "pytorch:22.03"] # 21.02, 21.10 for backward comp. + container: ["pytorch:21.02", "pytorch:21.10", "pytorch:22.06"] # 21.02, 21.10 for backward comp. container: image: nvcr.io/nvidia/${{ matrix.container }}-py3 # testing with the latest pytorch base image options: "--gpus all" runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: fetch-depth: 0 - name: Install the dependencies @@ -204,11 +204,11 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' needs: cron-gpu # so that monai itself is verified first container: - image: nvcr.io/nvidia/pytorch:22.03-py3 # testing with the latest pytorch base image + image: nvcr.io/nvidia/pytorch:22.06-py3 # testing with the latest pytorch base image options: "--gpus all --ipc=host" runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Install MONAI id: monai-install run: | @@ -223,7 +223,7 @@ jobs: - name: Checkout tutorials and install their requirements run: | cd /opt - git clone --depth 1 --branch master --single-branch https://github.com/Project-MONAI/tutorials.git # latest commit of master branch + git clone --depth 1 --branch main --single-branch https://github.com/Project-MONAI/tutorials.git # latest commit of main branch cd tutorials python -m pip install -r requirements.txt - name: Run tutorial notebooks diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 7140cd7dd8..4ca1b18261 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -19,13 +19,13 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 # full history so that we can git describe with: ref: dev fetch-depth: 0 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - shell: bash @@ -50,7 +50,7 @@ jobs: needs: versioning_dev runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: ref: dev - name: Download version diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index c38b66eda4..951b09d082 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -7,12 +7,12 @@ on: jobs: integration-py3: container: - image: nvcr.io/nvidia/pytorch:21.12-py3 # CUDA 11.5 - options: --gpus all + image: nvcr.io/nvidia/pytorch:22.04-py3 # CUDA 11.6 + options: --gpus all # shm-size 4g works fine runs-on: [self-hosted, linux, x64, common] steps: # checkout the pull request branch - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: token: ${{ secrets.PR_MAINTAIN }} repository: ${{ github.event.client_payload.pull_request.head.repo.full_name }} @@ -22,7 +22,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | @@ -34,7 +34,7 @@ jobs: which python python -m pip install --upgrade pip wheel python -m pip uninstall -y torch torchvision - python -m pip install torch==1.11.0+cu115 torchvision==0.12.0+cu115 -f https://download.pytorch.org/whl/torch_stable.html + python -m pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 -f https://download.pytorch.org/whl/torch_stable.html python -m pip install -r requirements-dev.txt rm -rf /github/home/.cache/torch/hub/mmars/ - name: Run integration tests diff --git a/.github/workflows/pythonapp-gpu.yml b/.github/workflows/pythonapp-gpu.yml index d19bbe437f..52e588a555 100644 --- a/.github/workflows/pythonapp-gpu.yml +++ b/.github/workflows/pythonapp-gpu.yml @@ -23,16 +23,17 @@ jobs: - "PT17+CUDA102" - "PT18+CUDA102" - "PT18+CUDA112" - - "PT111+CUDA116" + - "PT112+CUDA116" - "PT110+CUDA102" - - "PT111+CUDA102" + - "PT112+CUDA102" include: # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes - environment: PT17+CUDA102 pytorch: "torch==1.7.1 torchvision==0.8.2" base: "nvcr.io/nvidia/cuda:10.2-devel-ubuntu18.04" - environment: PT18+CUDA102 - pytorch: "torch==1.8.1 torchvision==0.9.1" + # pytorch 1.8.2 LTS + pytorch: "torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu102" base: "nvcr.io/nvidia/cuda:10.2-devel-ubuntu18.04" - environment: PT18+CUDA112 # we explicitly set pytorch to -h to avoid pip install error @@ -44,29 +45,36 @@ jobs: # 21.10: 1.10.0a0+0aef44c pytorch: "-h" base: "nvcr.io/nvidia/pytorch:21.10-py3" - - environment: PT111+CUDA116 + - environment: PT112+CUDA116 # we explicitly set pytorch to -h to avoid pip install error - # 22.03: 1.12.0a0+2c916ef + # 22.06: 1.13.0a0+340c412 pytorch: "-h" - base: "nvcr.io/nvidia/pytorch:22.03-py3" + base: "nvcr.io/nvidia/pytorch:22.06-py3" - environment: PT110+CUDA102 - pytorch: "torch==1.10.1 torchvision==0.11.2" + pytorch: "torch==1.10.2 torchvision==0.11.3" base: "nvcr.io/nvidia/cuda:10.2-devel-ubuntu18.04" - - environment: PT111+CUDA102 - pytorch: "torch==1.11.0 torchvision==0.12.0" + - environment: PT112+CUDA102 + pytorch: "torch==1.12.0 torchvision==0.13.0" base: "nvcr.io/nvidia/cuda:10.2-devel-ubuntu18.04" container: image: ${{ matrix.base }} options: --gpus all runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: apt install run: | + # workaround for https://github.com/Project-MONAI/MONAI/issues/4200 + apt-key del 7fa2af80 && rm -rf /etc/apt/sources.list.d/nvidia-ml.list /etc/apt/sources.list.d/cuda.list + apt-get update + apt-get install -y wget + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-keyring_1.0-1_all.deb + dpkg -i cuda-keyring_1.0-1_all.deb + if [ ${{ matrix.environment }} = "PT17+CUDA102" ] || \ [ ${{ matrix.environment }} = "PT18+CUDA102" ] || \ [ ${{ matrix.environment }} = "PT110+CUDA102" ] || \ - [ ${{ matrix.environment }} = "PT111+CUDA102" ] + [ ${{ matrix.environment }} = "PT112+CUDA102" ] then PYVER=3.7 PYSFX=3 DISTUTILS=python3-distutils && \ apt-get update && apt-get install -y --no-install-recommends \ diff --git a/.github/workflows/pythonapp-min.yml b/.github/workflows/pythonapp-min.yml index c3294c2b2a..8a9f4dbe67 100644 --- a/.github/workflows/pythonapp-min.yml +++ b/.github/workflows/pythonapp-min.yml @@ -27,9 +27,9 @@ jobs: os: [windows-latest, macOS-latest, ubuntu-latest] timeout-minutes: 40 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: Prepare pip wheel @@ -43,7 +43,7 @@ jobs: echo "::set-output name=dir::$(pip cache dir)" shell: bash - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ${{ steps.pip-cache.outputs.dir }} @@ -51,11 +51,11 @@ jobs: - if: runner.os == 'windows' name: Install torch cpu from pytorch.org (Windows only) run: | - python -m pip install torch==1.11.0+cpu -f https://download.pytorch.org/whl/torch_stable.html + python -m pip install torch==1.12.0+cpu -f https://download.pytorch.org/whl/torch_stable.html - name: Install the dependencies run: | # min. requirements - python -m pip install torch==1.11.0 + python -m pip install torch==1.12.0 python -m pip install -r requirements-min.txt python -m pip list BUILD_MONAI=0 python setup.py develop # no compile of extensions @@ -77,9 +77,9 @@ jobs: python-version: [3.7, 3.8, 3.9] timeout-minutes: 40 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - name: Prepare pip wheel @@ -93,7 +93,7 @@ jobs: echo "::set-output name=dir::$(pip cache dir)" shell: bash - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ${{ steps.pip-cache.outputs.dir }} @@ -101,7 +101,7 @@ jobs: - name: Install the dependencies run: | # min. requirements - python -m pip install torch==1.11.0 + python -m pip install torch==1.12.0 python -m pip install -r requirements-min.txt python -m pip list BUILD_MONAI=0 python setup.py develop # no compile of extensions @@ -119,12 +119,12 @@ jobs: strategy: fail-fast: false matrix: - pytorch-version: [1.6.0, 1.7.1, 1.8.1, 1.9.1, 1.10.1, latest] + pytorch-version: [1.7.1, 1.8.2, 1.9.1, 1.10.2, 1.11.0, latest] timeout-minutes: 40 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: Prepare pip wheel @@ -138,7 +138,7 @@ jobs: echo "::set-output name=dir::$(pip cache dir)" shell: bash - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ${{ steps.pip-cache.outputs.dir }} @@ -148,16 +148,16 @@ jobs: # min. requirements if [ ${{ matrix.pytorch-version }} == "latest" ]; then python -m pip install torch - elif [ ${{ matrix.pytorch-version }} == "1.6.0" ]; then - python -m pip install torch==1.6.0 elif [ ${{ matrix.pytorch-version }} == "1.7.1" ]; then python -m pip install torch==1.7.1 - elif [ ${{ matrix.pytorch-version }} == "1.8.1" ]; then - python -m pip install torch==1.8.1 + elif [ ${{ matrix.pytorch-version }} == "1.8.2" ]; then + python -m pip install torch==1.8.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cpu elif [ ${{ matrix.pytorch-version }} == "1.9.1" ]; then python -m pip install torch==1.9.1 - elif [ ${{ matrix.pytorch-version }} == "1.10.1" ]; then - python -m pip install torch==1.10.1 + elif [ ${{ matrix.pytorch-version }} == "1.10.2" ]; then + python -m pip install torch==1.10.2 + elif [ ${{ matrix.pytorch-version }} == "1.11.0" ]; then + python -m pip install torch==1.11.0 fi python -m pip install -r requirements-min.txt python -m pip list diff --git a/.github/workflows/pythonapp.yml b/.github/workflows/pythonapp.yml index cf251c2293..a0b1070bd4 100644 --- a/.github/workflows/pythonapp.yml +++ b/.github/workflows/pythonapp.yml @@ -22,9 +22,9 @@ jobs: flake8-py3: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: cache weekly timestamp @@ -32,7 +32,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ~/.cache/pip @@ -63,9 +63,9 @@ jobs: minimum-size: 8 maximum-size: 16 disk-root: "D:" - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: Prepare pip wheel @@ -79,7 +79,7 @@ jobs: echo "::set-output name=dir::$(pip cache dir)" shell: bash - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: ${{ steps.pip-cache.outputs.dir }} @@ -87,10 +87,10 @@ jobs: - if: runner.os == 'windows' name: Install torch cpu from pytorch.org (Windows only) run: | - python -m pip install torch==1.11.0+cpu torchvision==0.12.0+cpu -f https://download.pytorch.org/whl/torch_stable.html + python -m pip install torch==1.12.0+cpu torchvision==0.13.0+cpu -f https://download.pytorch.org/whl/torch_stable.html - name: Install the dependencies run: | - python -m pip install torch==1.11.0 torchvision==0.12.0 + python -m pip install torch==1.12.0 torchvision==0.13.0 cat "requirements-dev.txt" python -m pip install -r requirements-dev.txt python -m pip list @@ -105,6 +105,7 @@ jobs: python -m unittest -v env: QUICKTEST: True + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python # https://github.com/Project-MONAI/MONAI/issues/4354 packaging: runs-on: ubuntu-latest @@ -112,11 +113,11 @@ jobs: QUICKTEST: True shell: bash steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: fetch-depth: 0 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: cache weekly timestamp @@ -124,7 +125,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | @@ -137,7 +138,7 @@ jobs: # 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>=1.6 torchvision + python -m pip install torch>=1.7 torchvision - name: Check packages run: | pip uninstall monai @@ -189,13 +190,15 @@ jobs: ls -al python -m pip install -r requirements-dev.txt python -m unittest -v + env: + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python # https://github.com/Project-MONAI/MONAI/issues/4354 build-docs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: cache weekly timestamp @@ -203,7 +206,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 9cef1ff090..69f68d2119 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -15,11 +15,11 @@ jobs: matrix: python-version: [3.7, 3.8, 3.9] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - name: Install setuptools @@ -91,12 +91,12 @@ jobs: needs: packaging runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 # full history so that we can git describe with: fetch-depth: 0 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - shell: bash @@ -120,7 +120,7 @@ jobs: needs: versioning runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: Download version uses: actions/download-artifact@v2 with: diff --git a/.github/workflows/setupapp.yml b/.github/workflows/setupapp.yml index 02911ea51f..4a2a46cfc1 100644 --- a/.github/workflows/setupapp.yml +++ b/.github/workflows/setupapp.yml @@ -25,13 +25,13 @@ jobs: options: --gpus all runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 - name: cache weekly timestamp id: pip-cache run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | @@ -44,7 +44,7 @@ jobs: python -m pip install --upgrade pip wheel python -m pip uninstall -y torch torchvision rm -rf $(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")/ruamel* - python -m pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html + python -m pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html python -m pip install -r requirements-dev.txt - name: Run unit tests report coverage run: | @@ -76,11 +76,11 @@ jobs: matrix: python-version: [3.7, 3.8, 3.9] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - name: cache weekly timestamp @@ -88,7 +88,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | @@ -98,7 +98,7 @@ jobs: - name: Install the dependencies run: | python -m pip install --upgrade pip wheel - python -m pip install torch==1.11.0 torchvision==0.12.0 + python -m pip install torch==1.12.0 torchvision==0.13.0 python -m pip install -r requirements-dev.txt - name: Run quick tests CPU ubuntu run: | @@ -116,7 +116,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: cache weekly timestamp @@ -124,7 +124,7 @@ jobs: run: | echo "::set-output name=datew::$(date '+%Y-%V')" - name: cache for pip - uses: actions/cache@v2 + uses: actions/cache@v3 id: cache with: path: | @@ -146,7 +146,7 @@ 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@v2 + uses: actions/checkout@v3 with: ref: dev - name: Quick test installed (dev branch only) diff --git a/.github/workflows/weekly-preview.yml b/.github/workflows/weekly-preview.yml index 0d899b1f30..8fa69615c3 100644 --- a/.github/workflows/weekly-preview.yml +++ b/.github/workflows/weekly-preview.yml @@ -9,12 +9,12 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: ref: dev fetch-depth: 0 - name: Set up Python 3.8 - uses: actions/setup-python@v2 + uses: actions/setup-python@v3 with: python-version: 3.8 - name: Install setuptools @@ -33,7 +33,7 @@ jobs: 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 "0.9.dev${YEAR_WEEK}" + git tag "0.10.dev${YEAR_WEEK}" git log -1 git tag --list python setup.py sdist bdist_wheel diff --git a/.gitignore b/.gitignore index 542e08e3b6..9c0554dca9 100644 --- a/.gitignore +++ b/.gitignore @@ -131,6 +131,7 @@ tests/testing_data/MedNIST* tests/testing_data/*Hippocampus* tests/testing_data/*.tiff tests/testing_data/schema.json +*.svg # clang format tool .clang-format-bin/ @@ -138,3 +139,6 @@ tests/testing_data/schema.json # VSCode .vscode/ *.zip + +# profiling results +*.prof diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3fc762db46..d2a6990830 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ ci: repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.1.0 + rev: v4.3.0 hooks: - id: end-of-file-fixer - id: trailing-whitespace @@ -21,9 +21,14 @@ repos: - id: check-added-large-files args: ['--maxkb=1024'] - id: detect-private-key + - id: forbid-new-submodules + - id: pretty-format-json + args: ['--autofix', '--no-sort-keys', '--indent=4'] + - id: end-of-file-fixer + - id: mixed-line-ending - repo: https://github.com/asottile/pyupgrade - rev: v2.31.0 + rev: v2.34.0 hooks: - id: pyupgrade args: [--py37-plus] @@ -52,6 +57,12 @@ repos: docs/source/conf.py )$ + - repo: https://github.com/hadialqattan/pycln + rev: v1.3.5 + hooks: + - id: pycln + args: [--config=pyproject.toml] + #- repo: https://github.com/PyCQA/isort # rev: 5.9.3 # hooks: diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f55ded72f..bb0e8e31f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,57 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] +## [0.9.0] - 2022-06-08 +### Added +* `monai.bundle` primary module with a `ConfigParser` and command-line interfaces for configuration-based workflows +* Initial release of MONAI bundle specification +* Initial release of volumetric image detection modules including bounding boxes handling, RetinaNet-based architectures +* API preview `monai.data.MetaTensor` +* Unified `monai.data.image_writer` to support flexible IO backends including an ITK writer +* Various new network blocks and architectures including `SwinUNETR` +* DeepEdit interactive training/validation workflow +* NuClick interactive segmentation transforms +* Patch-based readers and datasets for whole-slide imaging +* New losses and metrics including `SurfaceDiceMetric`, `GeneralizedDiceFocalLoss` +* New pre-processing transforms including `RandIntensityRemap`, `SpatialResample` +* Multi-output and slice-based inference for `SlidingWindowInferer` +* `NrrdReader` for NRRD file support +* Torchscript utilities to save models with meta information +* Gradient-based visualization module `SmoothGrad` +* Automatic regular source code scanning for common vulnerabilities and coding errors + +### Changed +* Simplified `TestTimeAugmentation` using de-collate and invertible transforms APIs +* Refactoring `monai.apps.pathology` modules into `monai.handlers` and `monai.transforms` +* Flexible activation and normalization layers for `TopologySearch` and `DiNTS` +* Anisotropic first layers for 3D resnet +* Flexible ordering of activation, normalization in `UNet` +* Enhanced performance of connected-components analysis using Cupy +* `INSTANCE_NVFUSER` for enhanced performance in 3D instance norm +* Support of string representation of dtype in `convert_data_type` +* Added new options `iteration_log`, `iteration_log` to the logging handlers +* Base Docker image upgraded to `nvcr.io/nvidia/pytorch:22.04-py3` from `nvcr.io/nvidia/pytorch:21.10-py3` +* `collate_fn` generates more data-related debugging info with `dev_collate` + +### Fixed +* Unified the spellings of "meta data", "metadata", "meta-data" to "metadata" +* Various inaccurate error messages when input data are in invalid shapes +* Issue of computing symmetric distances in `compute_average_surface_distance` +* Unnecessary layer `self.conv3` in `UnetResBlock` +* Issue of torchscript compatibility for `ViT` and self-attention blocks +* Issue of hidden layers in `UNETR` +* `allow_smaller` in spatial cropping transforms +* Antialiasing in `Resize` +* Issue of bending energy loss value at different resolutions +* `kwargs_read_csv` in `CSVDataset` +* In-place modification in `Metric` reduction +* `wrap_array` for `ensure_tuple` +* Contribution guide for introducing new third-party dependencies + +### Removed +* Deprecated `nifti_writer`, `png_writer` in favor of `monai.data.image_writer` +* Support for PyTorch 1.6 + ## [0.8.1] - 2022-02-16 ### Added * Support of `matshow3d` with given `channel_dim` @@ -483,7 +534,8 @@ the postprocessing steps should be used before calling the metrics methods [highlights]: https://github.com/Project-MONAI/MONAI/blob/master/docs/source/highlights.md -[Unreleased]: https://github.com/Project-MONAI/MONAI/compare/0.8.1...HEAD +[Unreleased]: https://github.com/Project-MONAI/MONAI/compare/0.9.0...HEAD +[0.9.0]: https://github.com/Project-MONAI/MONAI/compare/0.8.1...0.9.0 [0.8.1]: https://github.com/Project-MONAI/MONAI/compare/0.8.0...0.8.1 [0.8.0]: https://github.com/Project-MONAI/MONAI/compare/0.7.0...0.8.0 [0.7.0]: https://github.com/Project-MONAI/MONAI/compare/0.6.0...0.7.0 diff --git a/CITATION.cff b/CITATION.cff index 8cbf686ce6..9fad6a709e 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -6,11 +6,15 @@ title: "MONAI: Medical Open Network for AI" abstract: "AI Toolkit for Healthcare Imaging" authors: - name: "MONAI Consortium" -date-released: 2020-03-28 -version: "0.6.0" -doi: "10.5281/zenodo.4323058" +date-released: 2022-06-13 +version: "0.9.0" +identifiers: + - description: "This DOI represents all versions of MONAI, and will always resolve to the latest one." + type: doi + value: "10.5281/zenodo.4323058" license: "Apache-2.0" repository-code: "https://github.com/Project-MONAI/MONAI" -cff-version: "1.1.0" +url: "https://monai.io" +cff-version: "1.2.0" message: "If you use this software, please cite it using these metadata." ... diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 129a839fd7..ac373bad75 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -5,6 +5,7 @@ 1. [Unit testing](#unit-testing) 1. [Building the documentation](#building-the-documentation) 1. [Automatic code formatting](#automatic-code-formatting) + 1. [Adding new optional dependencies](#adding-new-optional-dependencies) 1. [Signing your work](#signing-your-work) 1. [Utility functions](#utility-functions) 1. [Backwards compatibility](#backwards-compatibility) @@ -171,6 +172,48 @@ The first line of the comment must be `/black` so that it will be interpreted by - [Auto] After the formatting commit, the GitHub action adds an emoji to the comment that triggered the process. - Repeat the above steps if necessary. +#### Adding new optional dependencies +In addition to the minimal requirements of PyTorch and Numpy, MONAI's core modules are built optionally based on 3rd-party packages. +The current set of dependencies is listed in [installing dependencies](https://docs.monai.io/en/stable/installation.html#installing-the-recommended-dependencies). + +To allow for flexible integration of MONAI with other systems and environments, +the optional dependency APIs are always invoked lazily. For example, +```py +from monai.utils import optional_import +itk, _ = optional_import("itk", ...) + +class ITKReader(ImageReader): + ... + def read(self, ...): + return itk.imread(...) +``` +The availability of the external `itk.imread` API is not required unless `monai.data.ITKReader.read` is called by the user. +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) +- [docs/requirements.txt](https://github.com/Project-MONAI/MONAI/blob/dev/docs/requirements.txt) (pip requirements.txt 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 +an appropriate fallback default behaviour when the packages are not installed in +the testing environment. For example: +```py +from monai.utils import optional_import +plt, has_matplotlib = optional_import("matplotlib.pyplot") + +@skipUnless(has_matplotlib, "Matplotlib required") +class TestBlendImages(unittest.TestCase): +``` +It skips the test cases when `matplotlib.pyplot` APIs are not available. + +Alternatively, add the test file name to the ``exclude_cases`` in `tests/min_tests.py` to completely skip the test +cases when running in a minimal setup. + + + #### Signing your work MONAI enforces the [Developer Certificate of Origin](https://developercertificate.org/) (DCO) on all pull requests. All commit messages should contain the `Signed-off-by` line with an email address. The [GitHub DCO app](https://github.com/apps/dco) is deployed on MONAI. The pull request's status will be `failed` if commits do not contain a valid `Signed-off-by` line. diff --git a/Dockerfile b/Dockerfile index dc76584d5a..26a1fc242f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,7 +11,7 @@ # To build with a different base image # please run `docker build` using the `--build-arg PYTORCH_IMAGE=...` flag. -ARG PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:22.03-py3 +ARG PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:22.06-py3 FROM ${PYTORCH_IMAGE} LABEL maintainer="monai.contact@gmail.com" @@ -38,11 +38,11 @@ RUN BUILD_MONAI=1 FORCE_CUDA=1 python setup.py develop \ # NGC Client WORKDIR /opt/tools -ARG NGC_CLI_URI="https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip" -RUN wget -q ${NGC_CLI_URI} && \ - unzip ngccli_cat_linux.zip && chmod u+x ngc && \ - md5sum -c ngc.md5 && \ - rm -rf ngccli_cat_linux.zip ngc.md5 +ARG NGC_CLI_URI="https://ngc.nvidia.com/downloads/ngccli_linux.zip" +RUN wget -q ${NGC_CLI_URI} && unzip ngccli_linux.zip && chmod u+x ngc-cli/ngc && \ + find ngc-cli/ -type f -exec md5sum {} + | LC_ALL=C sort | md5sum -c ngc-cli.md5 && \ + rm -rf ngccli_linux.zip ngc-cli.md5 +ENV PATH=${PATH}:/opt/tools:/opt/tools/ngc-cli RUN apt-get update \ && DEBIAN_FRONTEND="noninteractive" apt-get install -y libopenslide0 \ && rm -rf /var/lib/apt/lists/* diff --git a/README.md b/README.md index bb217b7247..5701d1786e 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ To install [the current release](https://pypi.org/project/monai/), you can simpl pip install monai ``` -For other installation methods (using the default GitHub branch, using Docker, etc.), please refer to [the installation guide](https://docs.monai.io/en/latest/installation.html). +Please refer to [the installation guide](https://docs.monai.io/en/latest/installation.html) for other installation options. ## Getting Started @@ -56,12 +56,14 @@ Ask and answer questions over on [MONAI's GitHub Discussions tab](https://github ## Links - Website: https://monai.io/ -- API documentation: https://docs.monai.io +- API documentation (milestone): https://docs.monai.io/ +- API documentation (latest dev): https://docs.monai.io/en/latest/ - Code: https://github.com/Project-MONAI/MONAI - Project tracker: https://github.com/Project-MONAI/MONAI/projects - Issue tracker: https://github.com/Project-MONAI/MONAI/issues - Wiki: https://github.com/Project-MONAI/MONAI/wiki - Test status: https://github.com/Project-MONAI/MONAI/actions - PyPI package: https://pypi.org/project/monai/ +- conda-forge: https://anaconda.org/conda-forge/monai - Weekly previews: https://pypi.org/project/monai-weekly/ - Docker Hub: https://hub.docker.com/r/projectmonai/monai diff --git a/docs/Makefile b/docs/Makefile index 1a3b22d020..d9a870064a 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -8,6 +8,9 @@ SPHINXBUILD ?= sphinx-build SOURCEDIR = source BUILDDIR = build +# https://github.com/Project-MONAI/MONAI/issues/4354 +export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION := python + # Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/images/deepedit.png b/docs/images/deepedit.png new file mode 100644 index 0000000000..6da60db34e Binary files /dev/null and b/docs/images/deepedit.png differ diff --git a/docs/images/detection.png b/docs/images/detection.png new file mode 100644 index 0000000000..44637ced7f Binary files /dev/null and b/docs/images/detection.png differ diff --git a/docs/images/fast_training.png b/docs/images/fast_training.png index 34e47bcb21..8dbd1e5b8d 100644 Binary files a/docs/images/fast_training.png and b/docs/images/fast_training.png differ diff --git a/docs/images/nuclick.png b/docs/images/nuclick.png new file mode 100644 index 0000000000..b6768db5a2 Binary files /dev/null and b/docs/images/nuclick.png differ diff --git a/docs/images/swin_unetr.png b/docs/images/swin_unetr.png new file mode 100644 index 0000000000..82683a1715 Binary files /dev/null and b/docs/images/swin_unetr.png differ diff --git a/docs/requirements.txt b/docs/requirements.txt index f9749e9e36..09fff56b8b 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,6 +1,6 @@ -f https://download.pytorch.org/whl/cpu/torch-1.6.0%2Bcpu-cp37-cp37m-linux_x86_64.whl torch>=1.6 -pytorch-ignite==0.4.8 +pytorch-ignite==0.4.9 numpy>=1.17 itk>=5.2 nibabel @@ -28,3 +28,5 @@ tifffile; platform_system == "Linux" pyyaml fire jsonschema +pynrrd +pydicom diff --git a/docs/source/apps.rst b/docs/source/apps.rst index 239ae9eb17..3bb85c9296 100644 --- a/docs/source/apps.rst +++ b/docs/source/apps.rst @@ -122,3 +122,65 @@ Applications :members: .. autoclass:: TileOnGridd :members: + +`Detection` +----------- + +`Hard Negative Sampler` +~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.hard_negative_sampler + :members: + +`RetinaNet Network` +~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.networks.retinanet_network + :members: + +`RetinaNet Detector` +~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.networks.retinanet_detector + :members: + +`Transforms` +~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.transforms.box_ops + :members: +.. automodule:: monai.apps.detection.transforms.array + :members: +.. automodule:: monai.apps.detection.transforms.dictionary + :members: + +`Anchor` +~~~~~~~~ +.. automodule:: monai.apps.detection.utils.anchor_utils + :members: + +`Matcher` +~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.ATSS_matcher + :members: + +`Box coder` +~~~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.box_coder + :members: + +`Detection Utilities` +~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.detector_utils + :members: + +.. automodule:: monai.apps.detection.utils.predict_utils + :members: + +`Inference box selector` +~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.box_selector + :members: + +`Detection metrics` +~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.apps.detection.metrics.coco + :members: +.. automodule:: monai.apps.detection.metrics.matching + :members: diff --git a/docs/source/bundle.rst b/docs/source/bundle.rst index 297409cd7e..fc64d40f5a 100644 --- a/docs/source/bundle.rst +++ b/docs/source/bundle.rst @@ -36,6 +36,9 @@ Model Bundle `Scripts` --------- .. autofunction:: ckpt_export +.. autofunction:: download +.. autofunction:: load .. autofunction:: run .. autofunction:: verify_metadata .. autofunction:: verify_net_in_out +.. autofunction:: init_bundle diff --git a/docs/source/config_syntax.md b/docs/source/config_syntax.md index bcd849cd09..6cd4e62ef3 100644 --- a/docs/source/config_syntax.md +++ b/docs/source/config_syntax.md @@ -12,10 +12,10 @@ Content: - [A basic example](#a-basic-example) - [Syntax examples explained](#syntax-examples-explained) - - [`@` to interpolate with Python objects](#1--to-interpolate-with-python-objects) - - [`$` to evaluate as Python expressions](#2--to-evaluate-as-python-expressions) - - [`%` to textually replace configuration elements](#3--to-textually-replace-configuration-elements) - - [`_target_` (`_disabled_` and `_requires_`) to instantiate a Python object](#4-instantiate-a-python-object) + - [`@` to reference Python objects in configurations](#to-reference-python-objects-in-configurations) + - [`$` to evaluate as Python expressions](#to-evaluate-as-python-expressions) + - [`%` to textually replace configuration elements](#to-textually-replace-configuration-elements) + - [`_target_` (`_disabled_` and `_requires_`) to instantiate a Python object](#instantiate-a-python-object) - [The command line interface](#the-command-line-interface) - [Recommendations](#recommendations) @@ -73,22 +73,22 @@ For more details on the `ConfigParser` API, please see https://docs.monai.io/en/ A few characters and keywords are interpreted beyond the plain texts, here are examples of the syntax: -### 1. `@` to interpolate with Python objects +### To reference Python objects in configurations ```json "@preprocessing#transforms#keys" ``` -_Description:_ A reference to another configuration value defined at `preprocessing#transforms#keys`. +_Description:_ `@` character indicates a reference to another configuration value defined at `preprocessing#transforms#keys`. where `#` indicates a sub-structure of this configuration file. ```json "@preprocessing#1" ``` -_Description:_ `1` is interpreted as an integer, which is used to index (zero-based indexing) the `preprocessing` sub-structure. +_Description:_ `1` is referencing as an integer, which is used to index (zero-based indexing) the `preprocessing` sub-structure. -### 2. `$` to evaluate as Python expressions +### To evaluate as Python expressions ```json "$print(42)" @@ -110,16 +110,16 @@ _Description:_ `$` followed by an import statement is handled slightly different Python expressions. The imported module `resnet18` will be available as a global variable to the other configuration sections. This is to simplify the use of external modules in the configuration. -### 3. `%` to textually replace configuration elements +### To textually replace configuration elements ```json "%demo_config.json#demo_net#in_channels" ``` -_Description:_ A macro to replace the current configuration element with the texts at `demo_net#in_channels` in the +_Description:_ `%` character indicates a macro to replace the current configuration element with the texts at `demo_net#in_channels` in the `demo_config.json` file. The replacement is done before instantiating or evaluating the components. -### 4. instantiate a Python object +### Instantiate a Python object ```json { @@ -164,6 +164,7 @@ python -m monai.bundle COMMANDS where `COMMANDS` is one of the following: `run`, `verify_metadata`, `ckpt_export`, ... (please see `python -m monai.bundle --help` for a list of available options). +The CLI supports flexible use cases, such as overriding configs at runtime and predefining arguments in a file. To display a usage page for a command, for example `run`: ```bash python -m monai.bundle run -- --help @@ -182,3 +183,5 @@ Details on the CLI argument parsing is provided in the simple structures with sparse uses of expressions or references are preferred. - For `$import ` in the configuration, please make sure there are instructions for the users to install the `` if it is not a (optional) dependency of MONAI. +- As "#" and "$" might be interpreted differently by the `shell` or `CLI` tools, may need to add escape characters + or quotes for them in the command line, like: `"\$torch.device('cuda:1')"`, `"'train_part#trainer'"`. diff --git a/docs/source/data.rst b/docs/source/data.rst index 2bdf401c7f..eab4f867af 100644 --- a/docs/source/data.rst +++ b/docs/source/data.rst @@ -107,16 +107,23 @@ Patch-based dataset .. autoclass:: GridPatchDataset :members: -`PatchIter` -~~~~~~~~~~~ -.. autoclass:: PatchIter - :members: - `PatchDataset` ~~~~~~~~~~~~~~ .. autoclass:: PatchDataset :members: +`PatchIter` +""""""""""" +.. autoclass:: PatchIter + :members: + :special-members: __call__ + +`PatchIterd` +"""""""""""" +.. autoclass:: PatchIterd + :members: + :special-members: __call__ + Image reader ------------ @@ -145,10 +152,6 @@ PILReader .. autoclass:: PILReader :members: -WSIReader -~~~~~~~~~ -.. autoclass:: WSIReader - :members: Image writer ------------ @@ -264,3 +267,68 @@ ThreadDataLoader TestTimeAugmentation ~~~~~~~~~~~~~~~~~~~~ .. autoclass:: monai.data.TestTimeAugmentation + +N-Dim Fourier Transform +~~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.data.fft_utils +.. autofunction:: monai.data.fft_utils.fftn_centered +.. autofunction:: monai.data.fft_utils.ifftn_centered + + +Meta Object +----------- +.. automodule:: monai.data.meta_obj + :members: + +MetaTensor +---------- +.. autoclass:: monai.data.MetaTensor + :members: + + + +Whole slide image reader +------------------------ + +BaseWSIReader +~~~~~~~~~~~~~ +.. autoclass:: monai.data.BaseWSIReader + :members: + +WSIReader +~~~~~~~~~ +.. autoclass:: monai.data.WSIReader + :members: + +CuCIMWSIReader +~~~~~~~~~~~~~~ +.. autoclass:: monai.data.CuCIMWSIReader + :members: + +OpenSlideWSIReader +~~~~~~~~~~~~~~~~~~ +.. autoclass:: monai.data.OpenSlideWSIReader + :members: + +Whole slide image datasets +-------------------------- + +PatchWSIDataset +~~~~~~~~~~~~~~~ +.. autoclass:: monai.data.PatchWSIDataset + :members: + +MaskedPatchWSIDataset +~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: monai.data.MaskedPatchWSIDataset + :members: + +SlidingPatchWSIDataset +~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: monai.data.SlidingPatchWSIDataset + :members: + +Bounding box +------------ +.. automodule:: monai.data.box_utils + :members: diff --git a/docs/source/handlers.rst b/docs/source/handlers.rst index d32b6d88e3..0172529f40 100644 --- a/docs/source/handlers.rst +++ b/docs/source/handlers.rst @@ -95,12 +95,6 @@ Metric logger :members: -Segmentation saver ------------------- -.. autoclass:: SegmentationSaver - :members: - - Training stats handler ---------------------- .. autoclass:: StatsHandler @@ -174,3 +168,8 @@ Utilities --------- .. automodule:: monai.handlers.utils :members: + +Probability Map Handlers +------------------------ +.. automodule:: monai.handlers.probability_maps + :members: diff --git a/docs/source/highlights.md b/docs/source/highlights.md index cc91cfcd86..31e8eb21ee 100644 --- a/docs/source/highlights.md +++ b/docs/source/highlights.md @@ -24,6 +24,7 @@ The rest of this page provides more details for each module. * [Visualization](#visualization) * [Result writing](#result-writing) * [Workflows](#workflows) +* [Bundle](#bundle) * [Research](#research) * [Performance optimization and GPU acceleration](#performance-optimization-and-gpu-acceleration) * [Applications](#applications) @@ -328,11 +329,11 @@ To easily visualize a 3D image as frames of 2D images, MONAI provides the utilit ![matshow3d example](../images/matshow3d.png) -MONAI also provides the `blend_images` utility to blend the `image` and `label` to a RGB color image to better visualize the segmentation regions with the specified `cmap` mode and weights, etc. Showing a spleen segmentation `image` and the corresponding `label` as example: +MONAI also provides the `blend_images` utility to blend the `image` and `label` to an RGB color image to better visualize the segmentation regions with the specified `cmap` mode and weights, etc. Showing a spleen segmentation `image` and the corresponding `label` as example: ![blend example](../images/blend.png) -For more details of `TensorBoard utility`, `matshow3d` and `blend_images`, please check the [visualziation tutorial](https://github.com/Project-MONAI/tutorials/blob/master/modules/transform_visualization.ipynb). +For more details of `TensorBoard utility`, `matshow3d` and `blend_images`, please check the [visualization tutorial](https://github.com/Project-MONAI/tutorials/blob/master/modules/transform_visualization.ipynb). And to visualize the class activation mapping for a trained classification model, MONAI provides CAM, GradCAM, GradCAM++ APIs for both 2D and 3D models: @@ -343,7 +344,10 @@ The above example is generated by computing [GradCAM/GradCAM++ from a lung CT le ## Result writing Currently, MONAI supports writing the model outputs as NIfTI files or PNG files for segmentation tasks, and as CSV files for classification tasks. And the writers can restore the data spacing, orientation or shape according to the `original_shape` or `original_affine` information from the input image. -A rich set of formats will be supported soon, along with relevant statistics and evaluation metrics automatically computed from the outputs. +MONAI provides `SaveImage` transform to write the data volumes of image or prediction with automatically chose image writers based on the specified suffix, writers like: `NibabelWriter`, `ITKWriter`, `PILWriter`. +The `ImageWriter` API can be easily extended it for customized image writing modules. + +MONAI also supports to save the relevant statistics and evaluation metrics details automatically computed from the evaluation process. ## Workflows To quickly set up training and evaluation experiments, MONAI provides a set of workflows to significantly simplify the modules and allow for fast prototyping. @@ -382,7 +386,7 @@ The following figure compares the loss curves and validation scores for (1) trai ![transfer_mmar](../images/transfer_mmar.png) ### 5. Decollate batch data for flexible postprocessings -`decollate batch` is introduced in MONAI v0.6, which simplifies the post processing transforms and provides flexible following operations on a batch of data with various data shapes. It can decollate batched data (e.g. model predictions) into a list of tensors, for the benefits such as: +`decollate batch` is introduced in MONAI v0.6, which simplifies the post-processing transforms and provides flexible following operations on a batch of data with various data shapes. It can decollate batched data (e.g. model predictions) into a list of tensors, for the benefits such as: 1. enabling postprocessing transforms for each item independently -- randomised transforms could be applied differently for each predicted item in a batch. 2. simplifying the transform APIs and reducing the input validation burdens because both the preprocessing and postprocessing transforms now only need to support the "channel-first" input format. 3. enabling the `Invertd` transform for the predictions and the inverted data with different shapes, as the data items are in a list, not stacked in a single tensor. @@ -396,6 +400,34 @@ A typical process of `decollate batch` is illustrated as follows (with a `batch_ ### 6. Easy to integrate into popular workflows Except for the pytorch-ignite based `monai.engines`, most of the MONAI modules could be used independently or combined with other software packages. For example, MONAI can be easily integrated into popular frameworks such as PyTorch-Lightning and Catalyst: [Lightning segmentation](https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/spleen_segmentation_3d_lightning.ipynb) and [Lightning + TorchIO](https://github.com/Project-MONAI/tutorials/blob/master/modules/TorchIO_MONAI_PyTorch_Lightning.ipynb) tutorials show the PyTorch Lightning programs with MONAI modules, and [Catalyst segmentation](https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/unet_segmentation_3d_catalyst.ipynb) shows the Catalyst program with MONAI modules. +## Bundle +The objective of a MONAI bundle is to define a packaged model which includes the critical information necessary to allow users and programs to understand how the model is used and for what purpose. A bundle includes the stored weights of a single network as a pickled state dictionary plus optionally a Torchscript object and/or an ONNX object. Additional JSON files are included to store metadata about the model, information for constructing training, inference, and post-processing transform sequences, plain-text description, legal information, and other data the model creator wishes to include. More details are available at [bundle specification](https://docs.monai.io/en/latest/mb_specification.html). + +The key benefits of bundle are to define the model package and support building Python-based workflows via structured configurations: +- Self-contained model package include all the necessary information. +- Structured config can be used to easily reconstruct or prototype deep learning workflows. +- Config files can provide good readability and usability by separating parameter settings from the Python code. +- Config files can describe flexible workflow and components, allows for different low-level Python implementations +- Learning paradigms at a higher level such as federated learning and AutoML can be decoupled from the component details. + +A typical bundle example can include: +``` + ModelName + â”Ŗ━ configs + ┃ ┗━ metadata.json + â”Ŗ━ models + ┃ â”Ŗ━ model.pt + ┃ â”Ŗ━ *model.ts + ┃ ┗━ *model.onnx + ┗━ docs + â”Ŗ━ *README.md + ┗━ *license.txt +``` +Details about the bundle config definition and syntax & examples are at [config syntax](https://docs.monai.io/en/latest/config_syntax.html). +A step-by-step [get started](https://github.com/Project-MONAI/tutorials/blob/master/modules/bundles/get_started.ipynb) tutorial notebook can help users quickly set up a bundle. + +And [bundle examples](https://github.com/Project-MONAI/tutorials/tree/main/modules/bundle) provides more real-world examples and advanced features of bundle, including a bundle example for 3D segmentation of the spleen from CT image, use cases of bringing customized python components, parsing the config files in your own python program, etc. + ## Research There are several research prototypes in MONAI corresponding to the recently published papers that address advanced research problems. We always welcome contributions in forms of comments, suggestions, and code implementations. @@ -430,13 +462,20 @@ The [tutorial](https://github.com/Project-MONAI/tutorials/tree/master/self_super ![self-supervised](../images/ssl_overview.png) +### 6. Swin UNETR model for the task of multi-organ segmentation +For [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://arxiv.org/abs/2201.01266), MONAI introduces new network modules for multi-organ segmentation task using the BTCV challenge dataset. The architecture of Swin UNETR: + +![swin-unetr](../images/swin_unetr.png) + +The [tutorial](https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb) shows a typical pipeline of multi-organ segmentation based on Swin UNETR model, DiceCE loss function, Mean Dice, etc. And we used weights from self-supervised pre-training of Swin UNETR encoder (3D Swin Transformer) on a cohort of 5050 CT scans from publicly available datasets. + ## Performance optimization and GPU acceleration -Typically, model training is a time-consuming step during deep learning development, especially in medical imaging applications. Volumetric medical images are usually large (as multi-dimensional arrays) and the model training process can be complex. Even with powerful hardware (e.g. CPU/GPU with large RAM), it is not easy to fully leverage them to achieve high performance. MONAI provides [a fast training guide to achieve the best performance](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md). +Typically, model training is a time-consuming step during deep learning development, especially in medical imaging applications. Volumetric medical images are usually large (as multi-dimensional arrays) and the model training process can be complex. Even with powerful hardware (e.g. CPU/GPU with large RAM), it is not easy to fully leverage them to achieve high performance. MONAI provides a fast training guide to achieve the best performance: https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md. NVIDIA GPUs have been widely applied in many areas of deep learning training and evaluation, and the CUDA parallel computation shows obvious acceleration when comparing to traditional computation methods. To fully leverage GPU features, many popular mechanisms raised, like automatic mixed precision (AMP), distributed data parallel, etc. MONAI can support these features and provides rich examples. ### 1. Profiling the pipelines -First of all, MONAI provides several methods based on `DLProf`, `Nsight`, `NVTX` and `NVML` for users to analyze their programs to identify the performance bottleneck. The analyses include operation-based GPU activity and overall GPU activity during model training. They will greatly help users manage computing bottlenecks and provide insights for the area to be improved for better computing efficiency. The detailed example is shown in the [performance profiling tutorial](https://github.com/Project-MONAI/tutorials/blob/master/performance_profiling/radiology/profiling_train_base_nvtx.md). +First of all, MONAI provides several methods based on `DLProf`, `Nsight`, `NVTX` and `NVML` for users to analyze their programs to identify the performance bottleneck. The analyses include operation-based GPU activity and overall GPU activity during model training. They will greatly help users manage computing bottlenecks and provide insights for the area to be improved for better computing efficiency. The detailed example is shown in the performance profiling tutorial: https://github.com/Project-MONAI/tutorials/blob/master/performance_profiling/radiology/profiling_train_base_nvtx.md. ### 2. Auto mixed precision(AMP) In 2017, NVIDIA researchers developed a methodology for mixed-precision training, which combined single-precision (FP32) with half-precision (e.g. FP16) format when training a network, and it achieved the same accuracy as FP32 training using the same hyperparameters. @@ -448,7 +487,7 @@ MONAI workflows can easily set `amp=True/False` in `SupervisedTrainer` or `Super We also executed the same test program on NVIDIA A100 GPU with the same software environment, obtained faster results: ![amp a100 results](../images/amp_training_a100.png) More details is available at [AMP training tutorial](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/automatic_mixed_precision.ipynb). -We also tried to combine `AMP` with `CacheDataset`, `GPU cache`, `GPU transforms`, `ThreadDataLoader`, `DiceCE` loss function and `Novograd` optimizer to achieve the fast training in MONAI, able to obtain approximately `200x` speedup compared with a Pytorch native implementation when the training converges at a validation mean dice of `0.95`. Benchmark for reference: +We also tried to combine `AMP` with `CacheDataset`, `GPU cache`, `GPU transforms`, `ThreadDataLoader`, `DiceCE` loss, tuning of network and optimizer, to achieve the fast training in MONAI, with a V100 GPU and the target validation mean dice = 0.94 of the foreground channel only, it's more than `100x` speedup compared with the Pytorch regular implementation when achieving the same metric. And every epoch is 20x faster than regular training. Benchmark for reference: ![fast training results](../images/fast_training.png) More details is available at [Fast training tutorial](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_training_tutorial.ipynb). @@ -495,19 +534,40 @@ Sakinis, Tomas, et al. "Interactive segmentation of medical images through fully ![deepgrow scheme](../images/deepgrow.png) -### 2. Lesion detection in digital pathology -[Implementation](https://github.com/Project-MONAI/MONAI/tree/master/monai/apps/pathology) of the pathology detection components, which includes efficient whole slide imaging IO and sampling with NVIDIA cuCIM library and SmartCache mechanism, FROC measurements for lesion and probabilistic post-processing for lesion detection. +### 2. DeepEdit workflow for interactive segmentation +DeepEdit is a method that combines an automatic and a semi-automatic approach for 3D medical images into a single deep learning-based model. The [implementation](https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/deepedit) of the DeepEdit modules provides essential components for interactive segmentation. More details are available in the training and inference [tutorial](https://github.com/Project-MONAI/tutorials/tree/main/deepedit/ignite). + +The following figure shows the typical workflow of interactive segmentation: + +![deepedit workflow](../images/deepedit.png) + +### 3. NuClick modules for interactive nuclei segmentation +NuClick is a CNN-based approach to speed up collecting annotations for microscopic objects requiring minimum interaction from the annotator. The [implementation](https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/nuclick) contains essential components for the training and inference workflows of NuClick interactive nuclei segmentation. + +The following figure is example outputs of NuClick (annotator click inside the nucleus and the mask will be generated by CNN): + +![nuclick output](../images/nuclick.png) + +### 4. Lesion detection in digital pathology +[Implementation](https://github.com/Project-MONAI/MONAI/tree/master/monai/apps/pathology) of the pathology detection components, which includes efficient whole slide imaging IO and several patch sampling methods with NVIDIA cuCIM library and SmartCache mechanism, FROC measurements for lesion and probabilistic post-processing for lesion detection. ![digital pathology](../images/pathology.png) -### 3. Learning-based image registration +### 5. Learning-based image registration Starting from v0.5.0, MONAI provides experimental features for building learning-based 2D/3D registration workflows. These include image similarity measures as loss functions, bending energy as model regularization, network architectures, warping modules. The components can be used to build the major unsupervised and weakly-supervised algorithms. The following figure shows the registration of CT images acquired at different time points for a single patient using MONAI: ![3d registration](../images/3d_paired.png) -### 4. Reproducing the state-of-the-art Kaggle competition solutions +### 6. 2D and 3D detection workflow +The [implementation](https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/detection) contains 2D and 3D bounding box detection components of `RetinaNet`, which includesīŧšbounding box operations, hard negative sampler, and RetinaNet detectors. + +The following figure shows the detection training and inference workflows: + +![detection workflow](../images/detection.png) + +### 7. Reproducing the state-of-the-art Kaggle competition solutions [A reimplementation](https://github.com/Project-MONAI/tutorials/tree/master/kaggle/RANZCR/4th_place_solution) of the 4th place solution of RANZCR CLiP - Catheter and Line Position Challenge in Kaggle: https://www.kaggle.com/c/ranzcr-clip-catheter-line-classification The original solution is produced by Team Watercooled, and the authors are Dieter (https://www.kaggle.com/christofhenkel) and Psi (https://www.kaggle.com/philippsinger). diff --git a/docs/source/index.rst b/docs/source/index.rst index 6c99feac11..6bd8097ed7 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -76,7 +76,8 @@ Links ----- - Website: https://monai.io/ -- API documentation: https://docs.monai.io +- API documentation (milestone): https://docs.monai.io/ +- API documentation (latest dev): https://docs.monai.io/en/latest/ - Code: https://github.com/Project-MONAI/MONAI - Project tracker: https://github.com/Project-MONAI/MONAI/projects - Issue tracker: https://github.com/Project-MONAI/MONAI/issues @@ -85,6 +86,7 @@ Links - FAQ: https://github.com/Project-MONAI/MONAI/wiki/Frequently-asked-questions-and-answers - Test status: https://github.com/Project-MONAI/MONAI/actions - PyPI package: https://pypi.org/project/monai/ +- conda-forge: https://anaconda.org/conda-forge/monai - Weekly previews: https://pypi.org/project/monai-weekly/ - Docker Hub: https://hub.docker.com/r/projectmonai/monai diff --git a/docs/source/installation.md b/docs/source/installation.md index 12bf544cba..8623faee21 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -11,7 +11,7 @@ 3. [Validating the install](#validating-the-install) 4. [MONAI version string](#monai-version-string) 5. [From DockerHub](#from-dockerhub) -6. [Installing the recommended dependencies](#Installing-the-recommended-dependencies) +6. [Installing the recommended dependencies](#installing-the-recommended-dependencies) --- @@ -50,7 +50,7 @@ python -c "import monai; print(monai.__version__); print(monai.__commit_id__)" ## From conda-forge -To install the [current milestone release](https://pypi.org/project/monai/): +To install the [current milestone release](https://anaconda.org/conda-forge/monai): ```bash conda install -c conda-forge monai ``` @@ -190,9 +190,9 @@ Since MONAI v0.2.0, the extras syntax such as `pip install 'monai[nibabel]'` is - The options are ``` -[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema] +[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, pynrrd, pydicom] ``` which correspond to `nibabel`, `scikit-image`, `pillow`, `tensorboard`, -`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, respectively. +`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `pynrrd`, `pydicom`, respectively. - `pip install 'monai[all]'` installs all the optional dependencies. diff --git a/docs/source/losses.rst b/docs/source/losses.rst index dfd8ce2ddb..f05e4dc9ff 100644 --- a/docs/source/losses.rst +++ b/docs/source/losses.rst @@ -53,6 +53,11 @@ Segmentation Losses .. autoclass:: DiceFocalLoss :members: +`GeneralizedDiceFocalLoss` +~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: GeneralizedDiceFocalLoss + :members: + `FocalLoss` ~~~~~~~~~~~ .. autoclass:: FocalLoss diff --git a/docs/source/mb_specification.rst b/docs/source/mb_specification.rst index 5bdfa148e2..4d021d8f60 100644 --- a/docs/source/mb_specification.rst +++ b/docs/source/mb_specification.rst @@ -6,7 +6,7 @@ MONAI Bundle Specification Overview ======== -This is the specification for the MONAI Bundle (MB) format of portable described deep learning models. The objective of a MB is to define a packaged network or model which includes the critical information necessary to allow users and programs to understand how the model is used and for what purpose. A bundle includes the stored weights of a model as a pickled state dictionary and/or a Torchscript object. Additional JSON files are included to store metadata about the model, information for constructing training, inference, and post-processing transform sequences, plain-text description, legal information, and other data the model creator wishes to include. +This is the specification for the MONAI Bundle (MB) format of portable described deep learning models. The objective of a MB is to define a packaged network or model which includes the critical information necessary to allow users and programs to understand how the model is used and for what purpose. A bundle includes the stored weights of a single network as a pickled state dictionary plus optionally a Torchscript object and/or an ONNX object. Additional JSON files are included to store metadata about the model, information for constructing training, inference, and post-processing transform sequences, plain-text description, legal information, and other data the model creator wishes to include. This specification defines the directory structure a bundle must have and the necessary files it must contain. Additional files may be included and the directory packaged into a zip file or included as extra files directly in a Torchscript file. @@ -22,33 +22,42 @@ A MONAI Bundle is defined primarily as a directory with a set of specifically na ┃ ┗━ metadata.json â”Ŗ━ models ┃ â”Ŗ━ model.pt - ┃ ┗━ model.ts + ┃ â”Ŗ━ *model.ts + ┃ ┗━ *model.onnx ┗━ docs - â”Ŗ━ README.md - ┗━ license.txt + â”Ŗ━ *README.md + ┗━ *license.txt -These files mostly are required to be present with the given names for the directory to define a valid bundle: +The following files are **required** to be present with the given filenames for the directory to define a valid bundle: * **metadata.json**: metadata information in JSON format relating to the type of model, definition of input and output tensors, versions of the model and used software, and other information described below. * **model.pt**: the state dictionary of a saved model, the information to instantiate the model must be found in the metadata file. + +The following files are optional but must have these names in the directory given above: + * **model.ts**: the Torchscript saved model if the model is compatible with being saved correctly in this format. +* **model.onnx**: the ONNX model if the model is compatible with being saved correctly in this format. * **README.md**: plain-language information on the model, how to use it, author information, etc. in Markdown format. * **license.txt**: software license attached to the model, can be left blank if no license needed. +Other files can be included in any of the above directories. For example, `configs` can contain further configuration JSON or YAML files to define scripts for training or inference, overriding configuration values, environment definitions such as network instantiations, and so forth. One common file to include is `inference.json` which is used to define a basic inference script which uses input files with the stored network to produce prediction output files. + Archive Format ============== -The bundle directory and its contents can be compressed into a zip file to constitute a single file package. When unzipped into a directory this file will reproduce the above directory structure, and should itself also be named after the model it contains. +The bundle directory and its contents can be compressed into a zip file to constitute a single file package. When unzipped into a directory this file will reproduce the above directory structure, and should itself also be named after the model it contains. For example, `ModelName.zip` would contain at least `ModelName/configs/metadata.json` and `ModelName/models/model.pt`, thus when unzipped would place files into the directory `ModelName` rather than into the current working directory. + +The Torchscript file format is also just a zip file with a specific structure. When creating such an archive with `save_net_with_metadata` a MB-compliant Torchscript file can be created by including the contents of `metadata.json` as the `meta_values` argument of the function, and other files included as `more_extra_files` entries. These will be stored in a `extras` directory in the zip file and can be retrieved with `load_net_with_metadata` or with any other library/tool that can read zip data. In this format the `model.*` files are obviously not needed but `README.md` and `license.txt` as well as any others provided can be added as more extra files. -The Torchscript file format is also just a zip file with a specific structure. When creating such an archive with `save_net_with_metadata` a MB-compliant Torchscript file can be created by including the contents of `metadata.json` as the `meta_values` argument of the function, and other files included as `more_extra_files` entries. These will be stored in a `extras` directory in the zip file and can be retrieved with `load_net_with_metadata` or with any other library/tool that can read zip data. In this format the `model.*` files are obviously not needed by `README.md` and `license.txt` can be added as more extra files. +The `bundle` submodule of MONAI contains a number of command line programs. To produce a Torchscript bundle use `ckpt_export` with a set of specified components such as the saved weights file and metadata file. Config files can be provided as JSON or YAML dictionaries defining Python constructs used by the `ConfigParser`, however regardless of format the produced bundle Torchscript object will store the files as JSON. metadata.json File ================== This file contains the metadata information relating to the model, including what the shape and format of inputs and outputs are, what the meaning of the outputs are, what type of model is present, and other information. The JSON structure is a dictionary containing a defined set of keys with additional user-specified keys. The mandatory keys are as follows: -* **version**: version of the stored model, this allows multiple versions of the same model to be differentiated. +* **version**: version of the stored model, this allows multiple versions of the same model to be differentiated. Versions should follow semantic versioning and contain only characters valid in filenames as we may include the version to construct bundle file name. * **monai_version**: version of MONAI the bundle was generated on, later versions expected to work. * **pytorch_version**: version of Pytorch the bundle was generated on, later versions expected to work. * **numpy_version**: version of Numpy the bundle was generated on, later versions expected to work. @@ -61,8 +70,9 @@ This file contains the metadata information relating to the model, including wha Tensor format specifiers are used to define input and output tensors and their meanings, and must be a dictionary containing at least these keys: -* **type**: what sort of data the tensor represents: "image", "label", etc. -* **format**: what format of information is stored: "magnitude", "hounsfield", "kspace", "segmentation", "multiclass", etc. +* **type**: what sort of data the tensor represents: "image" for any spatial regular data whether an actual image or just data with that sort of shape, "series" for (time-) sequences of values such as signals, "tuples" for a series of items defined by a known number of values such as N-sized points in ND space, "probabilities" for a set of probabilities such as classifier output, this useful for interpreting what the dimensions and shape of the data represent and allow users to guess how to plot the data +* **format**: what format of information is stored, see below for list of known formats +* **modality**: describes the modality, protocol type, sort of capturing technology, or other property of the data not described by either it's type or format, known modalities are "MR", "CT", "US", "EKG", but can include any custom types or protocol types (eg. "T1"), default value is "n/a" * **num_channels**: number of channels the tensor has, assumed channel dimension first * **spatial_shape**: shape of the spatial dimensions of the form "[H]", "[H, W]", or "[H, W, D]", see below for possible values of H, W, and D * **dtype**: data type of tensor, eg. "float32", "int32" @@ -78,21 +88,38 @@ Optional keys: * **data_type**: type of source data used for training/validation. * **references**: list of published referenced relating to the model. -Spatial shape definition can be complex for models accepting inputs of varying shapes, especially if there are specific conditions on what those shapes can be. Shapes are specified as lists of either positive integers for fixed sizes or strings containing expressions defining the condition a size depends on. This can be "*" to mean any size, or use an expression with Python mathematical operators and one character variables to represent dependence on an unknown quantity. For example, "2**n" represents a size which must be a power of 2, "2**n*m" must be a multiple of a power of 2. Variables are shared between dimension expressions, so a spatial shape of `["2**n", "2**n"]` states that the dimensions must be the same powers of 2 given by `n`. +The format for tensors used as inputs and outputs can be used to specify semantic meaning of these values, and later is used by software handling bundles to determine how to process and interpret this data. There are various types of image data that MONAI is uses, and other data types such as point clouds, dictionary sequences, time signals, and others. The following list is provided as a set of supported definitions of what a tensor "format" is but is not exhaustive and users can provide their own which would be left up to the model users to interpret: -A JSON schema for this file can be found at https://github.com/Project-MONAI/MONAI/blob/3049e280f2424962bb2a69261389fcc0b98e0036/monai/apps/mmars/schema/metadata.json +* **magnitude**: ND field of continuous magnitude values with one or more channels, eg. MR T1 image having 1 channel or natural RGB image with 3 channels +* **hounsfield**: ND field of semi-categorical values given in Hounsfield, eg. CT image +* **kspace**: 2D/3D fourier transform image associated with MR imaging +* **raw**: ND field of values considered unprocessed from an image acquisition device, eg. directly from a MR scanner without reconstruction or other processing +* **labels**: ND categorical image with N one-hot channels for N-class segmentation/labels, the "channel_def" states in plain language what the interpretation of each channel is, for each pixel/voxel the predicted label is the index of the largest channel value +* **classes**: ND categorical image with N channels for N-class classes, the "channel_def" states in plain language what the interpretation of each channel is, this permits multi-class labeling as the channels need not be one-hot encoded +* **segmentation**: ND categorical image with one channel assigning each pixel/voxel to a label described in "channel_def" +* **points**: list of points/nodes/coordinates/vertices/vectors in ND space, so having a shape of [I, N] for I points with N dimensions +* **normals**: list of vectors (possible of unit length) in ND space, so having a shape of [I, N] for I vectors with N dimensions +* **indices**: list of indices into a vertices array and/or other array representing a set of shapes, so having a shape of [I, N] for I shapes defined by N values +* **sequence**: time-related sequence of values having one or more channels, such as a signal or dictionary lookup sentence, so having a shape of [C, N] for C channels of data at N time points. +* **latent**: ND tensor of data from the latent space from some layer of a network +* **gradient**: ND tensor of gradients from some layer of a network + +Spatial shape definition can be complex for models accepting inputs of varying shapes, especially if there are specific conditions on what those shapes can be. Shapes are specified as lists of either positive integers for fixed sizes or strings containing expressions defining the condition a size depends on. This can be "*" to mean any size, or use an expression with Python mathematical operators and one character variables to represent dependence on an unknown quantity. For example, "2**p" represents a size which must be a power of 2, "2**p*n" must be a multiple of a power of 2. Variables are shared between dimension expressions, a spatial shape example: `["*", "16*n", "2**p*n"]`. + +The download link of a JSON schema to verify this file can be found within it with key "schema". An example JSON metadata file: :: { + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json", "version": "0.1.0", "changelog": { "0.1.0": "complete the model package", "0.0.1": "initialize the model package structure" }, - "monai_version": "0.8.0", + "monai_version": "0.9.0", "pytorch_version": "1.10.0", "numpy_version": "1.21.2", "optional_packages_version": {"nibabel": "3.2.1"}, @@ -118,6 +145,7 @@ An example JSON metadata file: "image": { "type": "image", "format": "magnitude", + "modality": "MR", "num_channels": 1, "spatial_shape": [160, 160, 160], "dtype": "float32", @@ -129,11 +157,11 @@ An example JSON metadata file: "outputs":{ "pred": { "type": "image", - "format": "segmentation", + "format": "labels", "num_channels": 2, "spatial_shape": [160, 160, 160], "dtype": "float32", - "value_range": [0, 1], + "value_range": [], "is_patch_data": false, "channel_def": {0: "background", 1: "spleen"} } diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst index 332571d345..cb9cbe3c1d 100644 --- a/docs/source/metrics.rst +++ b/docs/source/metrics.rst @@ -39,6 +39,13 @@ Metrics .. autoclass:: DiceMetric :members: +`Generalized Dice Score` +------------------------ +.. autofunction:: compute_generalized_dice + +.. autoclass:: GeneralizedDiceScore + :members: + `Area under the ROC curve` -------------------------- .. autofunction:: compute_roc_auc @@ -69,6 +76,13 @@ Metrics .. autoclass:: SurfaceDistanceMetric :members: +`Surface dice` +-------------- +.. autofunction:: compute_surface_dice + +.. autoclass:: SurfaceDiceMetric + :members: + `Mean squared error` -------------------- .. autoclass:: MSEMetric diff --git a/docs/source/networks.rst b/docs/source/networks.rst index 7607cd2701..a7b6d8cdc0 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -40,6 +40,19 @@ Blocks .. autoclass:: MemoryEfficientSwish :members: +`FPN` +~~~~~ +.. autoclass:: ExtraFPNBlock + :members: +.. autoclass:: FeaturePyramidNetwork + :members: +.. autoclass:: LastLevelMaxPool + :members: +.. autoclass:: LastLevelP6P7 + :members: +.. autoclass:: BackboneWithFPN + :members: + `Mish` ~~~~~~ .. autoclass:: Mish @@ -220,6 +233,17 @@ Blocks .. autoclass:: DVF2DDF :members: + +N-Dim Fourier Transform +~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: monai.networks.blocks.fft_utils_t +.. autofunction:: monai.networks.blocks.fft_utils_t.fftn_centered_t +.. autofunction:: monai.networks.blocks.fft_utils_t.ifftn_centered_t +.. autofunction:: monai.networks.blocks.fft_utils_t.roll +.. autofunction:: monai.networks.blocks.fft_utils_t.roll_1d +.. autofunction:: monai.networks.blocks.fft_utils_t.fftshift +.. autofunction:: monai.networks.blocks.fft_utils_t.ifftshift + Layers ------ @@ -483,6 +507,11 @@ Nets .. autoclass:: UNETR :members: +`SwinUNETR` +~~~~~~~~~~~ +.. autoclass:: SwinUNETR + :members: + `BasicUNet` ~~~~~~~~~~~ .. autoclass:: BasicUNet @@ -580,11 +609,6 @@ Nets .. autoclass:: TorchVisionFCModel :members: -`TorchVisionFullyConvModel` -~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. autoclass:: TorchVisionFullyConvModel - :members: - `MILModel` ~~~~~~~~~~ .. autoclass:: MILModel diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst index 8fc832a253..7eaa17ea43 100644 --- a/docs/source/transforms.rst +++ b/docs/source/transforms.rst @@ -105,6 +105,12 @@ Crop and Pad :members: :special-members: __call__ +`Crop` +"""""" +.. autoclass:: Crop + :members: + :special-members: __call__ + `SpatialCrop` """"""""""""" .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/SpatialCrop.png @@ -446,6 +452,15 @@ Intensity :members: :special-members: __call__ + +`ForegroundMask` +"""""""""""""""" +.. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/ForegroundMask.png + :alt: example of ForegroundMask +.. autoclass:: ForegroundMask + :members: + :special-members: __call__ + IO ^^ @@ -737,6 +752,25 @@ Spatial :members: :special-members: __call__ +`GridPatch` +""""""""""" +.. autoclass:: GridPatch + :members: + :special-members: __call__ + +`RandGridPatch` +""""""""""""""" +.. autoclass:: RandGridPatch + :members: + :special-members: __call__ + +`GridSplit` +""""""""""" +.. autoclass:: GridSplit + :members: + :special-members: __call__ + + Smooth Field ^^^^^^^^^^^^ @@ -803,6 +837,12 @@ Utility :members: :special-members: __call__ +`SplitDim` +"""""""""" +.. autoclass:: SplitDim + :members: + :special-members: __call__ + `SplitChannel` """""""""""""" .. autoclass:: SplitChannel @@ -961,6 +1001,12 @@ Dictionary Transforms Crop and Pad (Dict) ^^^^^^^^^^^^^^^^^^^ +`Padd` +"""""" +.. autoclass:: Padd + :members: + :special-members: __call__ + `SpatialPadd` """"""""""""" .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/SpatialPadd.png @@ -985,6 +1031,18 @@ Crop and Pad (Dict) :members: :special-members: __call__ +`Cropd` +""""""" +.. autoclass:: Cropd + :members: + :special-members: __call__ + +`RandCropd` +""""""""""" +.. autoclass:: RandCropd + :members: + :special-members: __call__ + `SpatialCropd` """""""""""""" .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/SpatialCropd.png @@ -1314,6 +1372,13 @@ Intensity (Dict) :members: :special-members: __call__ +`ForegroundMaskd` +""""""""""""""""" +.. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/ForegroundMaskd.png + :alt: example of ForegroundMaskd +.. autoclass:: ForegroundMaskd + :members: + :special-members: __call__ IO (Dict) ^^^^^^^^^ @@ -1500,6 +1565,25 @@ Spatial (Dict) :members: :special-members: __call__ +`GridPatchd` +"""""""""""" +.. autoclass:: GridPatchd + :members: + :special-members: __call__ + +`RandGridPatchd` +"""""""""""""""" +.. autoclass:: RandGridPatchd + :members: + :special-members: __call__ + +`GridSplitd` +"""""""""""" +.. autoclass:: GridSplitd + :members: + :special-members: __call__ + + `RandRotate90d` """"""""""""""" .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/RandRotate90d.png @@ -1638,6 +1722,12 @@ Utility (Dict) :members: :special-members: __call__ +`SplitDimd` +""""""""""" +.. autoclass:: SplitDimd + :members: + :special-members: __call__ + `SplitChanneld` """"""""""""""" .. autoclass:: SplitChanneld @@ -1830,6 +1920,21 @@ Utility (Dict) :members: :special-members: __call__ +MetaTensor +^^^^^^^^^^ + +`ToMetaTensord` +""""""""""""""" +.. autoclass:: ToMetaTensord + :members: + :special-members: __call__ + +`FromMetaTensord` +""""""""""""""""" +.. autoclass:: FromMetaTensord + :members: + :special-members: __call__ + Transform Adaptors ------------------ .. automodule:: monai.transforms.adaptors diff --git a/docs/source/whatsnew.rst b/docs/source/whatsnew.rst index 1f27d651db..05c853b9cd 100644 --- a/docs/source/whatsnew.rst +++ b/docs/source/whatsnew.rst @@ -6,6 +6,7 @@ What's New .. toctree:: :maxdepth: 1 + whatsnew_0_9.md whatsnew_0_8.md whatsnew_0_7.md whatsnew_0_6.md diff --git a/docs/source/whatsnew_0_8.md b/docs/source/whatsnew_0_8.md index bbaf01f5de..3eb4bea167 100644 --- a/docs/source/whatsnew_0_8.md +++ b/docs/source/whatsnew_0_8.md @@ -1,4 +1,4 @@ -# What's new in 0.8 🎉🎉 +# What's new in 0.8 - Differentiable neural network topology search - Multiple instance learning for digital pathology WSI analysis diff --git a/docs/source/whatsnew_0_9.md b/docs/source/whatsnew_0_9.md new file mode 100644 index 0000000000..fa58630bdc --- /dev/null +++ b/docs/source/whatsnew_0_9.md @@ -0,0 +1,61 @@ +# What's new in 0.9 🎉🎉 + +- MONAI Bundle +- Object detection in medical images +- Swin Transformers for 3D medical image analysis +- New interactive segmentation components +- MetaTensor API preview + +## MONAI Bundle +MONAI Bundle format defines portable described of deep learning models ([docs](https://docs.monai.io/en/latest/bundle_intro.html)). +A bundle includes the critical information necessary during a model development life cycle, +and allows users and programs to understand the purpose and usage of the models. +The key benefits of Bundle and the `monai.bundle` APIs are: +- Standardized packaging format for storing and sharing models, +- Structured configuration files for fast prototyping of deep learning workflows, +- Easy to program APIs to separate deep learning hyperparameter settings from the Python code, +- Flexible config components to allow for different low-level Python implementations, +- Help to decouple the component details from higher level learning paradigms such as federated learning and AutoML. + +More details are [in the tutorials](https://github.com/Project-MONAI/tutorials/tree/main/modules/bundle). + +## Object detection in medical images +This release includes essential components for object localization and categorization workflows. +The initial developments include 2D and 3D bounding box handling, network blocks and architectures of RetinaNet, +and common utility modules such as coordinate-based preprocessing, hard negative sampler. + +The application specific modules are made available at +[monai.apps.detection](https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/detection). + +![detection workflow](../images/detection.png) + + +## Swin Transformers for 3D medical image analysis +The Swin UNETR model is now implemented in MONAI. +[The tutorial](https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb) +shows examples of multi-organ segmentation using this state-of-the-art model, +with weights from self-supervised pre-training of +Swin UNETR encoder (3D Swin Transformer) on a cohort of 5050 CT scans from publicly available datasets. +[The research-contribution entry](https://github.com/Project-MONAI/research-contributions/tree/main/SwinUNETR) +includes further technical details. + +![swin-unetr](../images/swin_unetr.png) + +## New interactive segmentation components +New components from deep learning interactive segmentation workflows +such as [DeepEdit](https://github.com/Project-MONAI/tutorials/tree/main/deepedit/ignite) +and NuClick are integrated into the core codebase. They serve as basic building blocks for +[the latest features in MONAILabel](https://github.com/Project-MONAI/MONAILabel). + +![deepedit](../images/deepedit.png) + +![nuclick](../images/nuclick.png) + +## MetaTensor API preview +The metadata associated with the primary imaging modalities is important in many biomedical applications, +especially for the data-driven approaches that MONAI has been focusing. +Starting from this release, we roll out a major refactoring for data representation in MONAI. For the first +step, [the core data structures](https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/meta_tensor.py) +`MetaTensor` and `MetaObj` are implemented as a feature preview. +Further developments [on the feature branch](https://github.com/Project-MONAI/MONAI/tree/feature/MetaTensor) +will be made available in future milestone releases. diff --git a/environment-dev.yml b/environment-dev.yml index a361262930..e71f434dc4 100644 --- a/environment-dev.yml +++ b/environment-dev.yml @@ -10,7 +10,7 @@ dependencies: - parameterized - setuptools>=50.3.0,!=60.0.0 - ignite==0.4.8 - - gdown>=3.6.4 + - gdown>=4.4.0 - scipy - nibabel - pillow!=8.3.0 # https://github.com/python-pillow/Pillow/issues/5571 @@ -45,6 +45,8 @@ dependencies: - pyyaml - fire - jsonschema + - pynrrd + - pydicom - pip - pip: # pip for itk as conda-forge version only up to v5.1 diff --git a/monai/_extensions/gmm/gmm.cpp b/monai/_extensions/gmm/gmm.cpp index 686fddb721..4087095340 100644 --- a/monai/_extensions/gmm/gmm.cpp +++ b/monai/_extensions/gmm/gmm.cpp @@ -15,75 +15,71 @@ limitations under the License. #include "gmm.h" -py::tuple init() -{ - torch::Tensor gmm_tensor = torch::zeros({GMM_COUNT, GMM_COMPONENT_COUNT}, torch::dtype(torch::kFloat32).device(torch::kCUDA)); - torch::Tensor scratch_tensor = torch::empty({1}, torch::dtype(torch::kFloat32).device(torch::kCUDA)); - return py::make_tuple(gmm_tensor, scratch_tensor); +py::tuple init() { + torch::Tensor gmm_tensor = + torch::zeros({GMM_COUNT, GMM_COMPONENT_COUNT}, torch::dtype(torch::kFloat32).device(torch::kCUDA)); + torch::Tensor scratch_tensor = torch::empty({1}, torch::dtype(torch::kFloat32).device(torch::kCUDA)); + return py::make_tuple(gmm_tensor, scratch_tensor); } -void learn(torch::Tensor gmm_tensor, torch::Tensor scratch_tensor, torch::Tensor input_tensor, torch::Tensor label_tensor) -{ - c10::DeviceType device_type = input_tensor.device().type(); - - unsigned int batch_count = input_tensor.size(0); - unsigned int element_count = input_tensor.stride(1); - - unsigned int scratch_size = batch_count * (element_count + GMM_COMPONENT_COUNT * GMM_COUNT * (element_count / (32 * 32))); - - if (scratch_tensor.size(0) < scratch_size) - { - scratch_tensor.resize_({scratch_size}); - } - - float* gmm = gmm_tensor.data_ptr(); - float* scratch = scratch_tensor.data_ptr(); - float* input = input_tensor.data_ptr(); - int* labels = label_tensor.data_ptr(); - - if(device_type == torch::kCUDA) - { - learn_cuda(input, labels, gmm, scratch, batch_count, element_count); - } - else - { - learn_cpu(input, labels, gmm, scratch, batch_count, element_count); - } +void learn( + torch::Tensor gmm_tensor, + torch::Tensor scratch_tensor, + torch::Tensor input_tensor, + torch::Tensor label_tensor) { + c10::DeviceType device_type = input_tensor.device().type(); + + unsigned int batch_count = input_tensor.size(0); + unsigned int element_count = input_tensor.stride(1); + + unsigned int scratch_size = + batch_count * (element_count + GMM_COMPONENT_COUNT * GMM_COUNT * (element_count / (32 * 32))); + + if (scratch_tensor.size(0) < scratch_size) { + scratch_tensor.resize_({scratch_size}); + } + + float* gmm = gmm_tensor.data_ptr(); + float* scratch = scratch_tensor.data_ptr(); + float* input = input_tensor.data_ptr(); + int* labels = label_tensor.data_ptr(); + + if (device_type == torch::kCUDA) { + learn_cuda(input, labels, gmm, scratch, batch_count, element_count); + } else { + learn_cpu(input, labels, gmm, scratch, batch_count, element_count); + } } -torch::Tensor apply(torch::Tensor gmm_tensor, torch::Tensor input_tensor) -{ - c10::DeviceType device_type = input_tensor.device().type(); - - unsigned int dim = input_tensor.dim(); - unsigned int batch_count = input_tensor.size(0); - unsigned int element_count = input_tensor.stride(1); - - long int* output_size = new long int[dim]; - memcpy(output_size, input_tensor.sizes().data(), dim * sizeof(long int)); - output_size[1] = MIXTURE_COUNT; - torch::Tensor output_tensor = torch::empty(c10::IntArrayRef(output_size, dim), torch::dtype(torch::kFloat32).device(device_type)); - delete output_size; - - const float* gmm = gmm_tensor.data_ptr(); - const float* input = input_tensor.data_ptr(); - float* output = output_tensor.data_ptr(); - - if(device_type == torch::kCUDA) - { - apply_cuda(gmm, input, output, batch_count, element_count); - } - else - { - apply_cpu(gmm, input, output, batch_count, element_count); - } - - return output_tensor; +torch::Tensor apply(torch::Tensor gmm_tensor, torch::Tensor input_tensor) { + c10::DeviceType device_type = input_tensor.device().type(); + + unsigned int dim = input_tensor.dim(); + unsigned int batch_count = input_tensor.size(0); + unsigned int element_count = input_tensor.stride(1); + + long int* output_size = new long int[dim]; + memcpy(output_size, input_tensor.sizes().data(), dim * sizeof(long int)); + output_size[1] = MIXTURE_COUNT; + torch::Tensor output_tensor = + torch::empty(c10::IntArrayRef(output_size, dim), torch::dtype(torch::kFloat32).device(device_type)); + delete output_size; + + const float* gmm = gmm_tensor.data_ptr(); + const float* input = input_tensor.data_ptr(); + float* output = output_tensor.data_ptr(); + + if (device_type == torch::kCUDA) { + apply_cuda(gmm, input, output, batch_count, element_count); + } else { + apply_cpu(gmm, input, output, batch_count, element_count); + } + + return output_tensor; } -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("init", torch::wrap_pybind_function(init)); - m.def("learn", torch::wrap_pybind_function(learn)); - m.def("apply", torch::wrap_pybind_function(apply)); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("init", torch::wrap_pybind_function(init)); + m.def("learn", torch::wrap_pybind_function(learn)); + m.def("apply", torch::wrap_pybind_function(apply)); } diff --git a/monai/_extensions/gmm/gmm.h b/monai/_extensions/gmm/gmm.h index 6317baa41a..09c0389ae6 100644 --- a/monai/_extensions/gmm/gmm.h +++ b/monai/_extensions/gmm/gmm.h @@ -24,9 +24,30 @@ limitations under the License. #define GMM_COMPONENT_COUNT (MATRIX_COMPONENT_COUNT + 1) #define GMM_COUNT (MIXTURE_COUNT * MIXTURE_SIZE) +void learn_cpu( + const float* input, + const int* labels, + float* gmm, + float* scratch_memory, + unsigned int batch_count, + unsigned int element_count); +void apply_cpu( + const float* gmm, + const float* input, + float* output, + unsigned int batch_count, + unsigned int element_count); -void learn_cpu(const float* input, const int* labels, float* gmm, float* scratch_memory, unsigned int batch_count, unsigned int element_count); -void apply_cpu(const float* gmm, const float* input, float* output, unsigned int batch_count, unsigned int element_count); - -void learn_cuda(const float* input, const int* labels, float* gmm, float* scratch_memory, unsigned int batch_count, unsigned int element_count); -void apply_cuda(const float* gmm, const float* input, float* output, unsigned int batch_count, unsigned int element_count); +void learn_cuda( + const float* input, + const int* labels, + float* gmm, + float* scratch_memory, + unsigned int batch_count, + unsigned int element_count); +void apply_cuda( + const float* gmm, + const float* input, + float* output, + unsigned int batch_count, + unsigned int element_count); diff --git a/monai/_extensions/gmm/gmm_cpu.cpp b/monai/_extensions/gmm/gmm_cpu.cpp index c9b55490eb..d7eedc07c8 100644 --- a/monai/_extensions/gmm/gmm_cpu.cpp +++ b/monai/_extensions/gmm/gmm_cpu.cpp @@ -15,12 +15,21 @@ limitations under the License. #include "gmm.h" -void learn_cpu(const float* input, const int* labels, float* gmm, float* scratch_memory, unsigned int batch_count, unsigned int element_count) -{ - throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu"); +void learn_cpu( + const float* input, + const int* labels, + float* gmm, + float* scratch_memory, + unsigned int batch_count, + unsigned int element_count) { + throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu"); } -void apply_cpu(const float* gmm, const float* input, float* output, unsigned int batch_count, unsigned int element_count) -{ - throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu"); +void apply_cpu( + const float* gmm, + const float* input, + float* output, + unsigned int batch_count, + unsigned int element_count) { + throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu"); } diff --git a/monai/_extensions/gmm/gmm_cuda.cu b/monai/_extensions/gmm/gmm_cuda.cu index 765ffe5b1c..2cf70a9920 100644 --- a/monai/_extensions/gmm/gmm_cuda.cu +++ b/monai/_extensions/gmm/gmm_cuda.cu @@ -20,434 +20,408 @@ limitations under the License. #define EPSILON 1e-5 #define BLOCK_SIZE 32 -#define TILE(SIZE, STRIDE) ((((SIZE) - 1)/(STRIDE)) + 1) +#define TILE(SIZE, STRIDE) ((((SIZE)-1) / (STRIDE)) + 1) -template -__global__ void CovarianceReductionKernel(int gaussian_index, const float* g_image, const int* g_alpha, float* g_matrices, int element_count) -{ - constexpr int block_size = warp_count * 32; +template +__global__ void CovarianceReductionKernel( + int gaussian_index, + const float* g_image, + const int* g_alpha, + float* g_matrices, + int element_count) { + constexpr int block_size = warp_count * 32; - __shared__ float s_matrix_component[warp_count]; + __shared__ float s_matrix_component[warp_count]; - int batch_index = blockIdx.z; + int batch_index = blockIdx.z; - const float* g_batch_image = g_image + batch_index * element_count * CHANNEL_COUNT; - const int* g_batch_alpha = g_alpha + batch_index * element_count; - float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * gridDim.x; + const float* g_batch_image = g_image + batch_index * element_count * CHANNEL_COUNT; + const int* g_batch_alpha = g_alpha + batch_index * element_count; + float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * gridDim.x; - int local_index = threadIdx.x; - int block_index = blockIdx.x; - int warp_index = local_index >> 5; - int lane_index = local_index & 31; - int global_index = local_index + block_index * block_size * load_count; - int matrix_offset = (gaussian_index * gridDim.x + block_index) * GMM_COMPONENT_COUNT; + int local_index = threadIdx.x; + int block_index = blockIdx.x; + int warp_index = local_index >> 5; + int lane_index = local_index & 31; + int global_index = local_index + block_index * block_size * load_count; + int matrix_offset = (gaussian_index * gridDim.x + block_index) * GMM_COMPONENT_COUNT; - float matrix[MATRIX_COMPONENT_COUNT]; + float matrix[MATRIX_COMPONENT_COUNT]; - for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) - { - matrix[i] = 0; - } + for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) { + matrix[i] = 0; + } + + for (int load = 0; load < load_count; load++) { + global_index += load * block_size; + + if (global_index < element_count) { + int my_alpha = g_batch_alpha[global_index]; + + if (my_alpha != -1) { + if (gaussian_index == (my_alpha & 15) + (my_alpha >> 4) * MIXTURE_COUNT) { + float feature[CHANNEL_COUNT + 1]; - for (int load = 0; load < load_count; load++) - { - global_index += load * block_size; - - if (global_index < element_count) - { - int my_alpha = g_batch_alpha[global_index]; - - if (my_alpha != -1) - { - if (gaussian_index == (my_alpha & 15) + (my_alpha >> 4) * MIXTURE_COUNT) - { - float feature[CHANNEL_COUNT + 1]; - - feature[0] = 1; - - for (int i = 0; i < CHANNEL_COUNT; i++) - { - feature[i + 1] = g_batch_image[global_index + i * element_count]; - } - - for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) - { - for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) - { - matrix[index] += feature[i] * feature[j]; - } - } - } + feature[0] = 1; + + for (int i = 0; i < CHANNEL_COUNT; i++) { + feature[i + 1] = g_batch_image[global_index + i * element_count]; + } + + for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) { + for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) { + matrix[index] += feature[i] * feature[j]; } + } } + } } + } - __syncthreads(); + __syncthreads(); - for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) - { - float matrix_component = matrix[i]; + for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) { + float matrix_component = matrix[i]; - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); - - if (lane_index == 0) - { - s_matrix_component[warp_index] = matrix_component; - } - - __syncthreads(); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); - if (warp_index == 0) - { - matrix_component = s_matrix_component[lane_index]; + if (lane_index == 0) { + s_matrix_component[warp_index] = matrix_component; + } - if (warp_count >= 32) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); } - if (warp_count >= 16) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); } - if (warp_count >= 8) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); } - if (warp_count >= 4) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); } - if (warp_count >= 2) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); } + __syncthreads(); - if (lane_index == 0) - { - g_batch_matrices[matrix_offset + i] = matrix_component; - } - } + if (warp_index == 0) { + matrix_component = s_matrix_component[lane_index]; - __syncthreads(); + if (warp_count >= 32) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); + } + if (warp_count >= 16) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); + } + if (warp_count >= 8) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); + } + if (warp_count >= 4) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); + } + if (warp_count >= 2) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); + } + + if (lane_index == 0) { + g_batch_matrices[matrix_offset + i] = matrix_component; + } } + + __syncthreads(); + } } -template -__global__ void CovarianceFinalizationKernel(const float* g_matrices, float* g_gmm, int matrix_count) -{ - constexpr int block_size = warp_count * 32; +template +__global__ void CovarianceFinalizationKernel(const float* g_matrices, float* g_gmm, int matrix_count) { + constexpr int block_size = warp_count * 32; - __shared__ float s_matrix_component[warp_count]; - __shared__ float s_gmm[GMM_COMPONENT_COUNT]; + __shared__ float s_matrix_component[warp_count]; + __shared__ float s_gmm[GMM_COMPONENT_COUNT]; - int batch_index = blockIdx.z; + int batch_index = blockIdx.z; - const float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * matrix_count; - float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; + const float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * matrix_count; + float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; - int local_index = threadIdx.x; - int warp_index = local_index >> 5; - int lane_index = local_index & 31; - int gmm_index = blockIdx.x; - int matrix_offset = gmm_index * matrix_count; + int local_index = threadIdx.x; + int warp_index = local_index >> 5; + int lane_index = local_index & 31; + int gmm_index = blockIdx.x; + int matrix_offset = gmm_index * matrix_count; - int load_count = TILE(matrix_count, block_size); + int load_count = TILE(matrix_count, block_size); - float norm_factor = 1.0f; + float norm_factor = 1.0f; - for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) - { - for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) - { - float matrix_component = 0.0f; + for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) { + for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) { + float matrix_component = 0.0f; - for(int load = 0; load < load_count; load++) - { - int matrix_index = local_index + load * block_size; + for (int load = 0; load < load_count; load++) { + int matrix_index = local_index + load * block_size; - if(matrix_index < matrix_count) - { - matrix_component += g_batch_matrices[(matrix_offset + matrix_index) * GMM_COMPONENT_COUNT + index]; - } - } + if (matrix_index < matrix_count) { + matrix_component += g_batch_matrices[(matrix_offset + matrix_index) * GMM_COMPONENT_COUNT + index]; + } + } - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); - matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); - if (lane_index == 0) - { - s_matrix_component[warp_index] = matrix_component; - } + if (lane_index == 0) { + s_matrix_component[warp_index] = matrix_component; + } - __syncthreads(); + __syncthreads(); - if (warp_index == 0) - { - matrix_component = s_matrix_component[lane_index]; + if (warp_index == 0) { + matrix_component = s_matrix_component[lane_index]; - if (warp_count >= 32) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); } - if (warp_count >= 16) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); } - if (warp_count >= 8) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); } - if (warp_count >= 4) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); } - if (warp_count >= 2) { matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); } - - if (lane_index == 0) - { - float constant = i == 0 ? 0.0f : s_gmm[i] * s_gmm[j]; + if (warp_count >= 32) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16); + } + if (warp_count >= 16) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8); + } + if (warp_count >= 8) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4); + } + if (warp_count >= 4) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2); + } + if (warp_count >= 2) { + matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1); + } - if (i != 0 && i == j) - { - constant -= EPSILON; - } + if (lane_index == 0) { + float constant = i == 0 ? 0.0f : s_gmm[i] * s_gmm[j]; - s_gmm[index] = norm_factor * matrix_component - constant; + if (i != 0 && i == j) { + constant -= EPSILON; + } - if (index == 0 && matrix_component > 0) - { - norm_factor = 1.0f / matrix_component; - } - } - } + s_gmm[index] = norm_factor * matrix_component - constant; - __syncthreads(); + if (index == 0 && matrix_component > 0) { + norm_factor = 1.0f / matrix_component; + } } + } + + __syncthreads(); } + } - float* matrix = s_gmm + (CHANNEL_COUNT + 1); - float* det_ptr = s_gmm + MATRIX_COMPONENT_COUNT; + float* matrix = s_gmm + (CHANNEL_COUNT + 1); + float* det_ptr = s_gmm + MATRIX_COMPONENT_COUNT; - if (local_index == 0) - { - float square_mat[CHANNEL_COUNT][CHANNEL_COUNT]; - float cholesky_mat[CHANNEL_COUNT][CHANNEL_COUNT]; + if (local_index == 0) { + float square_mat[CHANNEL_COUNT][CHANNEL_COUNT]; + float cholesky_mat[CHANNEL_COUNT][CHANNEL_COUNT]; - for(int i = 0; i < CHANNEL_COUNT; i++) - { - for(int j = 0; j < CHANNEL_COUNT; j++) - { - square_mat[i][j] = 0.0f; - cholesky_mat[i][j] = 0.0f; - } - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + for (int j = 0; j < CHANNEL_COUNT; j++) { + square_mat[i][j] = 0.0f; + cholesky_mat[i][j] = 0.0f; + } + } - to_square(matrix, square_mat); - cholesky(square_mat, cholesky_mat); + to_square(matrix, square_mat); + cholesky(square_mat, cholesky_mat); - *det_ptr = chol_det(cholesky_mat); + *det_ptr = chol_det(cholesky_mat); - if (invert_matrix) - { - chol_inv(cholesky_mat, square_mat); - to_triangle(square_mat, matrix); - } + if (invert_matrix) { + chol_inv(cholesky_mat, square_mat); + to_triangle(square_mat, matrix); } + } - if (local_index < GMM_COMPONENT_COUNT) - { - g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + local_index] = s_gmm[local_index]; - } + if (local_index < GMM_COMPONENT_COUNT) { + g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + local_index] = s_gmm[local_index]; + } } -struct GMMSplit_t -{ - int idx; - float threshold; - float eigenvector[CHANNEL_COUNT]; +struct GMMSplit_t { + int idx; + float threshold; + float eigenvector[CHANNEL_COUNT]; }; // 1 Block, 32xMIXTURE_COUNT -__global__ void GMMFindSplit(GMMSplit_t *gmmSplit, int gmmK, float *gmm) -{ - int batch_index = blockIdx.z; +__global__ void GMMFindSplit(GMMSplit_t* gmmSplit, int gmmK, float* gmm) { + int batch_index = blockIdx.z; - float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; - GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT; + float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; + GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT; - int gmm_idx = threadIdx.x * MIXTURE_COUNT + threadIdx.y; + int gmm_idx = threadIdx.x * MIXTURE_COUNT + threadIdx.y; - float eigenvalue = 0; - float eigenvector[CHANNEL_COUNT]; + float eigenvalue = 0; + float eigenvector[CHANNEL_COUNT]; - if (threadIdx.x < gmmK) - { - float* matrix = g_batch_gmm + gmm_idx * GMM_COMPONENT_COUNT + (CHANNEL_COUNT + 1); - largest_eigenpair(matrix, eigenvector, &eigenvalue); - } - - float max_value = eigenvalue; + if (threadIdx.x < gmmK) { + float* matrix = g_batch_gmm + gmm_idx * GMM_COMPONENT_COUNT + (CHANNEL_COUNT + 1); + largest_eigenpair(matrix, eigenvector, &eigenvalue); + } - max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 16)); - max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 8)); - max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 4)); - max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 2)); - max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 1)); + float max_value = eigenvalue; - if (max_value == eigenvalue) - { - GMMSplit_t split; + max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 16)); + max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 8)); + max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 4)); + max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 2)); + max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 1)); - float* average_feature = gmm + gmm_idx * GMM_COMPONENT_COUNT + 1; + if (max_value == eigenvalue) { + GMMSplit_t split; - split.idx = threadIdx.x; - split.threshold = scalar_prod(average_feature, eigenvector); + float* average_feature = gmm + gmm_idx * GMM_COMPONENT_COUNT + 1; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - split.eigenvector[i] = eigenvector[i]; - } + split.idx = threadIdx.x; + split.threshold = scalar_prod(average_feature, eigenvector); - g_batch_gmmSplit[threadIdx.y] = split; + for (int i = 0; i < CHANNEL_COUNT; i++) { + split.eigenvector[i] = eigenvector[i]; } + + g_batch_gmmSplit[threadIdx.y] = split; + } } #define DO_SPLIT_DEGENERACY 4 -__global__ void GMMDoSplit(const GMMSplit_t *gmmSplit, int k, const float *image, int *alpha, int element_count) -{ - __shared__ GMMSplit_t s_gmmSplit[MIXTURE_COUNT]; +__global__ void GMMDoSplit(const GMMSplit_t* gmmSplit, int k, const float* image, int* alpha, int element_count) { + __shared__ GMMSplit_t s_gmmSplit[MIXTURE_COUNT]; - int batch_index = blockIdx.z; + int batch_index = blockIdx.z; - const GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT; - const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT; - int* g_batch_alpha = alpha + batch_index * element_count; + const GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT; + const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT; + int* g_batch_alpha = alpha + batch_index * element_count; - int *s_linear = (int *) s_gmmSplit; - int *g_linear = (int *) g_batch_gmmSplit; + int* s_linear = (int*)s_gmmSplit; + int* g_linear = (int*)g_batch_gmmSplit; - if (threadIdx.x < MIXTURE_COUNT * sizeof(GMMSplit_t)) - { - s_linear[threadIdx.x] = g_linear[threadIdx.x]; - } + if (threadIdx.x < MIXTURE_COUNT * sizeof(GMMSplit_t)) { + s_linear[threadIdx.x] = g_linear[threadIdx.x]; + } - __syncthreads(); + __syncthreads(); - int index = threadIdx.x + blockIdx.x * BLOCK_SIZE * DO_SPLIT_DEGENERACY; + int index = threadIdx.x + blockIdx.x * BLOCK_SIZE * DO_SPLIT_DEGENERACY; - for (int i = 0; i < DO_SPLIT_DEGENERACY; i++) - { - index += BLOCK_SIZE; + for (int i = 0; i < DO_SPLIT_DEGENERACY; i++) { + index += BLOCK_SIZE; - if (index < element_count) - { - int my_alpha = g_batch_alpha[index]; + if (index < element_count) { + int my_alpha = g_batch_alpha[index]; - if(my_alpha != -1) - { - int select = my_alpha & 15; - int gmm_idx = my_alpha >> 4; + if (my_alpha != -1) { + int select = my_alpha & 15; + int gmm_idx = my_alpha >> 4; - if (gmm_idx == s_gmmSplit[select].idx) - { - // in the split cluster now - float feature[CHANNEL_COUNT]; + if (gmm_idx == s_gmmSplit[select].idx) { + // in the split cluster now + float feature[CHANNEL_COUNT]; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - feature[i] = g_batch_image[index + i * element_count]; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + feature[i] = g_batch_image[index + i * element_count]; + } - float value = scalar_prod(s_gmmSplit[select].eigenvector, feature); + float value = scalar_prod(s_gmmSplit[select].eigenvector, feature); - if (value > s_gmmSplit[select].threshold) - { - // assign pixel to new cluster - g_batch_alpha[index] = k + select; - } - } - } + if (value > s_gmmSplit[select].threshold) { + // assign pixel to new cluster + g_batch_alpha[index] = k + select; + } } + } } + } } // Single block, 32xMIXTURE_COUNT -__global__ void GMMcommonTerm(float *g_gmm) -{ - int batch_index = blockIdx.z; +__global__ void GMMcommonTerm(float* g_gmm) { + int batch_index = blockIdx.z; - float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; + float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; - int gmm_index = (threadIdx.x * MIXTURE_COUNT) + threadIdx.y; + int gmm_index = (threadIdx.x * MIXTURE_COUNT) + threadIdx.y; - float gmm_n = threadIdx.x < MIXTURE_SIZE ? g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT] : 0.0f; + float gmm_n = threadIdx.x < MIXTURE_SIZE ? g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT] : 0.0f; - float sum = gmm_n; + float sum = gmm_n; - sum += __shfl_xor_sync(0xffffffff, sum, 1); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 16); - if (threadIdx.x < MIXTURE_SIZE) - { - float det = g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] + EPSILON; - float commonTerm = det > 0.0f ? gmm_n / (sqrtf(det) * sum) : gmm_n / sum; + if (threadIdx.x < MIXTURE_SIZE) { + float det = g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] + EPSILON; + float commonTerm = det > 0.0f ? gmm_n / (sqrtf(det) * sum) : gmm_n / sum; - g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] = commonTerm; - } + g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] = commonTerm; + } } -__device__ float GMMTerm(float* feature, const float *gmm) -{ - const float* average_feature = gmm + 1; - const float* matrix = gmm + CHANNEL_COUNT + 1; +__device__ float GMMTerm(float* feature, const float* gmm) { + const float* average_feature = gmm + 1; + const float* matrix = gmm + CHANNEL_COUNT + 1; - float diff[CHANNEL_COUNT]; + float diff[CHANNEL_COUNT]; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - diff[i] = feature[i] - average_feature[i]; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + diff[i] = feature[i] - average_feature[i]; + } - float value = 0.0f; + float value = 0.0f; - for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) - { - for (int j = i; j < CHANNEL_COUNT; j++, index++) - { - float term = diff[i] * diff[j] * matrix[index]; + for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) { + for (int j = i; j < CHANNEL_COUNT; j++, index++) { + float term = diff[i] * diff[j] * matrix[index]; - value += i == j ? term : 2 * term; - } + value += i == j ? term : 2 * term; } + } - return gmm[MATRIX_COMPONENT_COUNT] * expf(-0.5 * value); + return gmm[MATRIX_COMPONENT_COUNT] * expf(-0.5 * value); } -__global__ void GMMDataTermKernel(const float *image, const float *gmm, float* output, int element_count) -{ - int batch_index = blockIdx.z; +__global__ void GMMDataTermKernel(const float* image, const float* gmm, float* output, int element_count) { + int batch_index = blockIdx.z; - const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT; - const float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; - float* g_batch_output = output + batch_index * element_count * MIXTURE_COUNT; + const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT; + const float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT; + float* g_batch_output = output + batch_index * element_count * MIXTURE_COUNT; - int index = blockIdx.x * blockDim.x + threadIdx.x; + int index = blockIdx.x * blockDim.x + threadIdx.x; - if (index >= element_count) return; + if (index >= element_count) + return; - float feature[CHANNEL_COUNT]; + float feature[CHANNEL_COUNT]; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - feature[i] = g_batch_image[index + i * element_count]; - } - - float weights[MIXTURE_COUNT]; - float weight_total = 0.0f; + for (int i = 0; i < CHANNEL_COUNT; i++) { + feature[i] = g_batch_image[index + i * element_count]; + } - for(int i = 0; i < MIXTURE_COUNT; i++) - { - float mixture_weight = 0.0f; + float weights[MIXTURE_COUNT]; + float weight_total = 0.0f; - for(int j = 0; j < MIXTURE_SIZE; j++) - { - mixture_weight += GMMTerm(feature, &g_batch_gmm[(MIXTURE_COUNT * j + i) * GMM_COMPONENT_COUNT]); - } + for (int i = 0; i < MIXTURE_COUNT; i++) { + float mixture_weight = 0.0f; - weights[i] = mixture_weight; - weight_total += mixture_weight; + for (int j = 0; j < MIXTURE_SIZE; j++) { + mixture_weight += GMMTerm(feature, &g_batch_gmm[(MIXTURE_COUNT * j + i) * GMM_COMPONENT_COUNT]); } - for(int i = 0; i < MIXTURE_COUNT; i++) - { - // protecting against pixels with 0 in all mixtures - float final_weight = weight_total > 0.0f ? weights[i] / weight_total : 0.0f; - g_batch_output[index + i * element_count] = final_weight; - } + weights[i] = mixture_weight; + weight_total += mixture_weight; + } + + for (int i = 0; i < MIXTURE_COUNT; i++) { + // protecting against pixels with 0 in all mixtures + float final_weight = weight_total > 0.0f ? weights[i] / weight_total : 0.0f; + g_batch_output[index + i * element_count] = final_weight; + } } #define THREADS 512 @@ -455,67 +429,90 @@ __global__ void GMMDataTermKernel(const float *image, const float *gmm, float* o #define BLOCK (WARPS << 5) #define LOAD 4 -void GMMInitialize(const float *image, int *alpha, float *gmm, float *scratch_mem, unsigned int batch_count, unsigned int element_count) -{ - unsigned int block_count = TILE(element_count, BLOCK * LOAD); +void GMMInitialize( + const float* image, + int* alpha, + float* gmm, + float* scratch_mem, + unsigned int batch_count, + unsigned int element_count) { + unsigned int block_count = TILE(element_count, BLOCK * LOAD); - float* block_gmm_scratch = scratch_mem; - GMMSplit_t* gmm_split_scratch = (GMMSplit_t*) scratch_mem; + float* block_gmm_scratch = scratch_mem; + GMMSplit_t* gmm_split_scratch = (GMMSplit_t*)scratch_mem; - int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE; + int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE; - for (unsigned int k = MIXTURE_COUNT; k < gmm_N; k+=MIXTURE_COUNT) - { - for (unsigned int i = 0; i < k; ++i) - { - CovarianceReductionKernel<<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count); - } + for (unsigned int k = MIXTURE_COUNT; k < gmm_N; k += MIXTURE_COUNT) { + for (unsigned int i = 0; i < k; ++i) { + CovarianceReductionKernel + <<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count); + } - CovarianceFinalizationKernel<<<{k, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count); + CovarianceFinalizationKernel<<<{k, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count); - GMMFindSplit<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm_split_scratch, k / MIXTURE_COUNT, gmm); - GMMDoSplit<<<{TILE(element_count, BLOCK_SIZE * DO_SPLIT_DEGENERACY), 1, batch_count}, BLOCK_SIZE>>>(gmm_split_scratch, (k / MIXTURE_COUNT) << 4, image, alpha, element_count); - } + GMMFindSplit<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm_split_scratch, k / MIXTURE_COUNT, gmm); + GMMDoSplit<<<{TILE(element_count, BLOCK_SIZE * DO_SPLIT_DEGENERACY), 1, batch_count}, BLOCK_SIZE>>>( + gmm_split_scratch, (k / MIXTURE_COUNT) << 4, image, alpha, element_count); + } } -void GMMUpdate(const float *image, int *alpha, float *gmm, float *scratch_mem, unsigned int batch_count, unsigned int element_count) -{ - unsigned int block_count = TILE(element_count, BLOCK * LOAD); +void GMMUpdate( + const float* image, + int* alpha, + float* gmm, + float* scratch_mem, + unsigned int batch_count, + unsigned int element_count) { + unsigned int block_count = TILE(element_count, BLOCK * LOAD); - float* block_gmm_scratch = scratch_mem; + float* block_gmm_scratch = scratch_mem; - unsigned int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE; + unsigned int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE; - for (unsigned int i = 0; i < gmm_N; ++i) - { - CovarianceReductionKernel<<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count); - } + for (unsigned int i = 0; i < gmm_N; ++i) { + CovarianceReductionKernel + <<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count); + } - CovarianceFinalizationKernel<<<{gmm_N, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count); + CovarianceFinalizationKernel<<<{gmm_N, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count); - GMMcommonTerm<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm); + GMMcommonTerm<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm); } -void GMMDataTerm(const float *image, const float *gmm, float* output, unsigned int batch_count, unsigned int element_count) -{ - dim3 block(BLOCK_SIZE, 1); - dim3 grid(TILE(element_count, BLOCK_SIZE), 1, batch_count); +void GMMDataTerm( + const float* image, + const float* gmm, + float* output, + unsigned int batch_count, + unsigned int element_count) { + dim3 block(BLOCK_SIZE, 1); + dim3 grid(TILE(element_count, BLOCK_SIZE), 1, batch_count); - GMMDataTermKernel<<>>(image, gmm, output, element_count); + GMMDataTermKernel<<>>(image, gmm, output, element_count); } -void learn_cuda(const float* input, const int* labels, float* gmm, float* scratch_memory, unsigned int batch_count, unsigned int element_count) -{ - int* alpha = (int*)scratch_memory; - float* scratch_mem = scratch_memory + batch_count * element_count; +void learn_cuda( + const float* input, + const int* labels, + float* gmm, + float* scratch_memory, + unsigned int batch_count, + unsigned int element_count) { + int* alpha = (int*)scratch_memory; + float* scratch_mem = scratch_memory + batch_count * element_count; - cudaMemcpyAsync(alpha, labels, batch_count * element_count * sizeof(int), cudaMemcpyDeviceToDevice); + cudaMemcpyAsync(alpha, labels, batch_count * element_count * sizeof(int), cudaMemcpyDeviceToDevice); - GMMInitialize(input, alpha, gmm, scratch_mem, batch_count, element_count); - GMMUpdate(input, alpha, gmm, scratch_mem, batch_count, element_count); + GMMInitialize(input, alpha, gmm, scratch_mem, batch_count, element_count); + GMMUpdate(input, alpha, gmm, scratch_mem, batch_count, element_count); } -void apply_cuda(const float* gmm, const float* input, float* output, unsigned int batch_count, unsigned int element_count) -{ - GMMDataTerm(input, gmm, output, batch_count, element_count); +void apply_cuda( + const float* gmm, + const float* input, + float* output, + unsigned int batch_count, + unsigned int element_count) { + GMMDataTerm(input, gmm, output, batch_count, element_count); } diff --git a/monai/_extensions/gmm/gmm_cuda_linalg.cuh b/monai/_extensions/gmm/gmm_cuda_linalg.cuh index 9d54d80d3b..56c7c7ccdc 100644 --- a/monai/_extensions/gmm/gmm_cuda_linalg.cuh +++ b/monai/_extensions/gmm/gmm_cuda_linalg.cuh @@ -11,170 +11,134 @@ See the License for the specific language governing permissions and limitations under the License. */ -__device__ void to_square(float in[SUB_MATRIX_COMPONENT_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) -{ - for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) - { - for (int j = i; j < CHANNEL_COUNT; j++, index++) - { - out[i][j] = in[index]; - out[j][i] = in[index]; - } +__device__ void to_square(float in[SUB_MATRIX_COMPONENT_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) { + for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) { + for (int j = i; j < CHANNEL_COUNT; j++, index++) { + out[i][j] = in[index]; + out[j][i] = in[index]; } + } } -__device__ void to_triangle(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[SUB_MATRIX_COMPONENT_COUNT]) -{ - for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) - { - for (int j = i; j < CHANNEL_COUNT; j++, index++) - { - out[index] = in[j][i]; - } +__device__ void to_triangle(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[SUB_MATRIX_COMPONENT_COUNT]) { + for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) { + for (int j = i; j < CHANNEL_COUNT; j++, index++) { + out[index] = in[j][i]; } + } } -__device__ void cholesky(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) -{ - for (int i = 0; i < CHANNEL_COUNT; i++) - { - for (int j = 0; j < i+1; j++) - { - float sum = 0.0f; - - for (int k = 0; k < j; k++) - { - sum += out[i][k] * out[j][k]; - } - - if (i == j) - { - out[i][j] = sqrtf(in[i][i] - sum); - } - else - { - out[i][j] = (in[i][j] - sum) / out[j][j]; - } - } +__device__ void cholesky(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) { + for (int i = 0; i < CHANNEL_COUNT; i++) { + for (int j = 0; j < i + 1; j++) { + float sum = 0.0f; + + for (int k = 0; k < j; k++) { + sum += out[i][k] * out[j][k]; + } + + if (i == j) { + out[i][j] = sqrtf(in[i][i] - sum); + } else { + out[i][j] = (in[i][j] - sum) / out[j][j]; + } } + } } -__device__ float chol_det(float in[CHANNEL_COUNT][CHANNEL_COUNT]) -{ - float det = 1.0f; +__device__ float chol_det(float in[CHANNEL_COUNT][CHANNEL_COUNT]) { + float det = 1.0f; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - det *= in[i][i]; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + det *= in[i][i]; + } - return det * det; + return det * det; } -__device__ void chol_inv(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) -{ - // Invert cholesky matrix - for (int i = 0; i < CHANNEL_COUNT; i++) - { - in[i][i] = 1.0f / (in[i][i] + 0.0001f); +__device__ void chol_inv(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) { + // Invert cholesky matrix + for (int i = 0; i < CHANNEL_COUNT; i++) { + in[i][i] = 1.0f / (in[i][i] + 0.0001f); - for (int j = 0; j < i; j++) - { - float sum = 0.0f; + for (int j = 0; j < i; j++) { + float sum = 0.0f; - for (int k = j; k < i; k++) - { - sum += in[i][k] * in[k][j]; - } + for (int k = j; k < i; k++) { + sum += in[i][k] * in[k][j]; + } - in[i][j] = -in[i][i] * sum; - } + in[i][j] = -in[i][i] * sum; } + } - // Dot with transpose of self - for (int i = 0; i < CHANNEL_COUNT; i++) - { - for (int j = 0; j < CHANNEL_COUNT; j++) - { - out[i][j] = 0.0f; - - for (int k = max(i, j); k < CHANNEL_COUNT; k++) - { - out[i][j] += in[k][i] * in[k][j]; - } - } + // Dot with transpose of self + for (int i = 0; i < CHANNEL_COUNT; i++) { + for (int j = 0; j < CHANNEL_COUNT; j++) { + out[i][j] = 0.0f; + + for (int k = max(i, j); k < CHANNEL_COUNT; k++) { + out[i][j] += in[k][i] * in[k][j]; + } } + } } -__device__ void normalize(float* v) -{ - float norm = 0.0f; +__device__ void normalize(float* v) { + float norm = 0.0f; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - norm += v[i] * v[i]; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + norm += v[i] * v[i]; + } - norm = 1.0f / sqrtf(norm); + norm = 1.0f / sqrtf(norm); - for (int i = 0; i < CHANNEL_COUNT; i++) - { - v[i] *= norm; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + v[i] *= norm; + } } -__device__ float scalar_prod(float* a, float* b) -{ - float product = 0.0f; +__device__ float scalar_prod(float* a, float* b) { + float product = 0.0f; - for (int i = 0; i < CHANNEL_COUNT; i++) - { - product += a[i] * b[i]; - } + for (int i = 0; i < CHANNEL_COUNT; i++) { + product += a[i] * b[i]; + } - return product; + return product; } -__device__ void largest_eigenpair(const float *M, float* evec, float* eval) -{ - float scratch[CHANNEL_COUNT]; - - for(int i = 0; i < CHANNEL_COUNT; i++) - { - scratch[i] = i + 1; - } +__device__ void largest_eigenpair(const float* M, float* evec, float* eval) { + float scratch[CHANNEL_COUNT]; - for (int itr = 0; itr < 10; itr++) - { - *eval = 0.0f; + for (int i = 0; i < CHANNEL_COUNT; i++) { + scratch[i] = i + 1; + } - for (int i = 0; i < CHANNEL_COUNT; i++) - { - int index = i; + for (int itr = 0; itr < 10; itr++) { + *eval = 0.0f; - evec[i] = 0.0f; + for (int i = 0; i < CHANNEL_COUNT; i++) { + int index = i; - for (int j = 0; j < CHANNEL_COUNT; j++) - { - evec[i] += M[index] * scratch[j]; + evec[i] = 0.0f; - if (j < i) - { - index += CHANNEL_COUNT - (j + 1); - } - else - { - index += 1; - } - } + for (int j = 0; j < CHANNEL_COUNT; j++) { + evec[i] += M[index] * scratch[j]; - *eval = max(*eval, evec[i]); + if (j < i) { + index += CHANNEL_COUNT - (j + 1); + } else { + index += 1; } + } - for (int i = 0; i < CHANNEL_COUNT; i++) - { - evec[i] /= *eval; - scratch[i] = evec[i]; - } + *eval = max(*eval, evec[i]); + } + + for (int i = 0; i < CHANNEL_COUNT; i++) { + evec[i] /= *eval; + scratch[i] = evec[i]; } + } } diff --git a/monai/apps/deepedit/interaction.py b/monai/apps/deepedit/interaction.py new file mode 100644 index 0000000000..04fabec06d --- /dev/null +++ b/monai/apps/deepedit/interaction.py @@ -0,0 +1,100 @@ +# 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 typing import Callable, Dict, Sequence, Union + +import numpy as np +import torch + +from monai.data import decollate_batch, list_data_collate +from monai.engines import SupervisedEvaluator, SupervisedTrainer +from monai.engines.utils import IterationEvents +from monai.transforms import Compose +from monai.utils.enums import CommonKeys + + +class Interaction: + """ + Ignite process_function used to introduce interactions (simulation of clicks) for DeepEdit Training/Evaluation. + + More details about this can be found at: + + Diaz-Pinto et al., MONAI Label: A framework for AI-assisted Interactive + Labeling of 3D Medical Images. (2022) https://arxiv.org/abs/2203.12362 + + Args: + deepgrow_probability: probability of simulating clicks in an iteration + transforms: execute additional transformation during every iteration (before train). + Typically, several Tensor based transforms composed by `Compose`. + train: True for training mode or False for evaluation mode + click_probability_key: key to click/interaction probability + label_names: Dict of label names + """ + + def __init__( + self, + deepgrow_probability: float, + transforms: Union[Sequence[Callable], Callable], + train: bool, + label_names: Dict[str, int], + click_probability_key: str = "probability", + ) -> None: + + self.deepgrow_probability = deepgrow_probability + self.transforms = Compose(transforms) if not isinstance(transforms, Compose) else transforms + self.train = train + self.label_names = label_names + self.click_probability_key = click_probability_key + + def __call__(self, engine: Union[SupervisedTrainer, SupervisedEvaluator], batchdata: Dict[str, torch.Tensor]): + + if batchdata is None: + raise ValueError("Must provide batch data for current iteration.") + + if np.random.choice([True, False], p=[self.deepgrow_probability, 1 - self.deepgrow_probability]): + + # Run the inner loop only once + inputs, _ = engine.prepare_batch(batchdata) + inputs = inputs.to(engine.state.device) + + engine.fire_event(IterationEvents.INNER_ITERATION_STARTED) + + engine.network.eval() + with torch.no_grad(): + if engine.amp: + with torch.cuda.amp.autocast(): + predictions = engine.inferer(inputs, engine.network) + else: + predictions = engine.inferer(inputs, engine.network) + batchdata.update({CommonKeys.PRED: predictions}) + + # decollate/collate batchdata to execute click transforms + batchdata_list = decollate_batch(batchdata, detach=True) + + for i in range(len(batchdata_list)): + batchdata_list[i][self.click_probability_key] = 1.0 + batchdata_list[i] = self.transforms(batchdata_list[i]) + + batchdata = list_data_collate(batchdata_list) + + engine.fire_event(IterationEvents.INNER_ITERATION_COMPLETED) + + else: + # zero out input guidance channels + batchdata_list = decollate_batch(batchdata, detach=True) + for i in range(1, len(batchdata_list[0][CommonKeys.IMAGE])): + batchdata_list[0][CommonKeys.IMAGE][i] *= 0 + batchdata = list_data_collate(batchdata_list) + + # first item in batch only + engine.state.batch = batchdata + + return engine._iteration(engine, batchdata) diff --git a/monai/apps/deepedit/transforms.py b/monai/apps/deepedit/transforms.py index 0fcf4c2286..5a7068291a 100644 --- a/monai/apps/deepedit/transforms.py +++ b/monai/apps/deepedit/transforms.py @@ -11,14 +11,19 @@ import json import logging -from typing import Dict, Hashable, Mapping, Tuple +import random +import warnings +from typing import Dict, Hashable, Mapping, Optional import numpy as np +import torch from monai.config import KeysCollection +from monai.networks.layers import GaussianFilter from monai.transforms.transform import MapTransform, Randomizable, Transform -from monai.utils import optional_import -from monai.utils.enums import PostFix +from monai.utils import min_version, optional_import + +measure, _ = optional_import("skimage.measure", "0.14.2", min_version) logger = logging.getLogger(__name__) @@ -27,23 +32,39 @@ class DiscardAddGuidanced(MapTransform): - def __init__(self, keys: KeysCollection, probability: float = 1.0, allow_missing_keys: bool = False): + def __init__( + self, + keys: KeysCollection, + number_intensity_ch: int = 1, + probability: float = 1.0, + label_names=None, + allow_missing_keys: bool = False, + ): """ - Discard positive and negative points randomly or Add the two channels for inference time + Discard positive and negative points according to discard probability - :param probability: Discard probability; For inference it will be always 1.0 + Args: + keys: The ``keys`` parameter will be used to get and set the actual data item to transform + number_intensity_ch: number of intensity channels + probability: probability of discarding clicks """ super().__init__(keys, allow_missing_keys) - self.probability = probability + + self.number_intensity_ch = number_intensity_ch + self.discard_probability = probability + self.label_names = label_names def _apply(self, image): - if self.probability >= 1.0 or np.random.choice([True, False], p=[self.probability, 1 - self.probability]): - signal = np.zeros((1, image.shape[-3], image.shape[-2], image.shape[-1]), dtype=np.float32) - if image.shape[0] == 3: - image[1] = signal - image[2] = signal + if self.discard_probability >= 1.0 or np.random.choice( + [True, False], p=[self.discard_probability, 1 - self.discard_probability] + ): + signal = np.zeros( + (len(self.label_names), image.shape[-3], image.shape[-2], image.shape[-1]), dtype=np.float32 + ) + if image.shape[0] == self.number_intensity_ch + len(self.label_names): + image[self.number_intensity_ch :, ...] = signal else: - image = np.concatenate((image, signal, signal), axis=0) + image = np.concatenate([image, signal], axis=0) return image def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: @@ -56,51 +77,428 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda return d -class ResizeGuidanceCustomd(Transform): +class NormalizeLabelsInDatasetd(MapTransform): + def __init__(self, keys: KeysCollection, label_names=None, allow_missing_keys: bool = False): + """ + Normalize label values according to label names dictionary + + Args: + keys: The ``keys`` parameter will be used to get and set the actual data item to transform + label_names: all label names + """ + super().__init__(keys, allow_missing_keys) + + self.label_names = label_names + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + # Dictionary containing new label numbers + new_label_names = {} + label = np.zeros(d[key].shape) + # Making sure the range values and number of labels are the same + for idx, (key_label, val_label) in enumerate(self.label_names.items(), start=1): + if key_label != "background": + new_label_names[key_label] = idx + label[d[key] == val_label] = idx + if key_label == "background": + new_label_names["background"] = 0 + + d["label_names"] = new_label_names + d[key] = label + else: + warnings.warn("This transform only applies to the label") + return d + + +class SingleLabelSelectiond(MapTransform): + def __init__(self, keys: KeysCollection, label_names=None, allow_missing_keys: bool = False): + """ + Selects one label at a time to train the DeepEdit + + Args: + keys: The ``keys`` parameter will be used to get and set the actual data item to transform + label_names: all label names + """ + super().__init__(keys, allow_missing_keys) + + self.label_names = label_names + self.all_label_values = { + "spleen": 1, + "right kidney": 2, + "left kidney": 3, + "gallbladder": 4, + "esophagus": 5, + "liver": 6, + "stomach": 7, + "aorta": 8, + "inferior vena cava": 9, + "portal_vein": 10, + "splenic_vein": 11, + "pancreas": 12, + "right adrenal gland": 13, + "left adrenal gland": 14, + } + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + # Taking one label at a time + t_label = np.random.choice(self.label_names) + d["current_label"] = t_label + d[key][d[key] != self.all_label_values[t_label]] = 0.0 + # Convert label to index values following label_names argument + max_label_val = self.label_names.index(t_label) + 1 + d[key][d[key] > 0] = max_label_val + print(f"Using label {t_label} with number: {d[key].max()}") + else: + warnings.warn("This transform only applies to the label") + return d + + +class AddGuidanceSignalDeepEditd(MapTransform): """ - Resize the guidance based on cropped vs resized image. + Add Guidance signal for input image. Multilabel DeepEdit + + Based on the "guidance" points, apply Gaussian to them and add them as new channel for input image. + + Args: + guidance: key to store guidance. + sigma: standard deviation for Gaussian kernel. + number_intensity_ch: channel index. """ - def __init__(self, guidance: str, ref_image: str) -> None: + def __init__( + self, + keys: KeysCollection, + guidance: str = "guidance", + sigma: int = 3, + number_intensity_ch: int = 1, + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) self.guidance = guidance - self.ref_image = ref_image + self.sigma = sigma + self.number_intensity_ch = number_intensity_ch - def __call__(self, data): - d = dict(data) - current_shape = d[self.ref_image].shape[1:] + def _get_signal(self, image, guidance): + dimensions = 3 if len(image.shape) > 3 else 2 + guidance = guidance.tolist() if isinstance(guidance, np.ndarray) else guidance + guidance = json.loads(guidance) if isinstance(guidance, str) else guidance - factor = np.divide(current_shape, d[PostFix.meta("image")]["dim"][1:4]) - pos_clicks, neg_clicks = d["foreground"], d["background"] + # In inference the user may not provide clicks for some channels/labels + if len(guidance): + if dimensions == 3: + # Assume channel is first and depth is last CHWD + signal = np.zeros((1, image.shape[-3], image.shape[-2], image.shape[-1]), dtype=np.float32) + else: + signal = np.zeros((1, image.shape[-2], image.shape[-1]), dtype=np.float32) + + sshape = signal.shape + for point in guidance: # TO DO: make the guidance a list only - it is currently a list of list + if np.any(np.asarray(point) < 0): + continue + + if dimensions == 3: + # Making sure points fall inside the image dimension + p1 = max(0, min(int(point[-3]), sshape[-3] - 1)) + p2 = max(0, min(int(point[-2]), sshape[-2] - 1)) + p3 = max(0, min(int(point[-1]), sshape[-1] - 1)) + signal[:, p1, p2, p3] = 1.0 + else: + p1 = max(0, min(int(point[-2]), sshape[-2] - 1)) + p2 = max(0, min(int(point[-1]), sshape[-1] - 1)) + signal[:, p1, p2] = 1.0 + + # Apply a Gaussian filter to the signal + if np.max(signal[0]) > 0: + signal_tensor = torch.tensor(signal[0]) + pt_gaussian = GaussianFilter(len(signal_tensor.shape), sigma=self.sigma) + signal_tensor = pt_gaussian(signal_tensor.unsqueeze(0).unsqueeze(0)) + signal_tensor = signal_tensor.squeeze(0).squeeze(0) + signal[0] = signal_tensor.detach().cpu().numpy() + signal[0] = (signal[0] - np.min(signal[0])) / (np.max(signal[0]) - np.min(signal[0])) + return signal + else: + if dimensions == 3: + signal = np.zeros((1, image.shape[-3], image.shape[-2], image.shape[-1]), dtype=np.float32) + else: + signal = np.zeros((1, image.shape[-2], image.shape[-1]), dtype=np.float32) + return signal + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "image": + image = d[key] + tmp_image = image[0 : 0 + self.number_intensity_ch, ...] + guidance = d[self.guidance] + for key_label in guidance.keys(): + # Getting signal based on guidance + signal = self._get_signal(image, guidance[key_label]) + tmp_image = np.concatenate([tmp_image, signal], axis=0) + d[key] = tmp_image + return d + else: + print("This transform only applies to image key") + return d + + +class FindAllValidSlicesDeepEditd(MapTransform): + """ + Find/List all valid slices in the labels. + Label is assumed to be a 4D Volume with shape CHWD, where C=1. + + Args: + sids: key to store slices indices having valid label map. + """ + + def __init__(self, keys: KeysCollection, sids="sids", allow_missing_keys: bool = False): + super().__init__(keys, allow_missing_keys) + self.sids = sids + + def _apply(self, label, d): + sids = {} + for key_label in d["label_names"].keys(): + l_ids = [] + for sid in range(label.shape[-1]): # Assume channel is first and depth is last CHWD + if d["label_names"][key_label] in label[0][..., sid]: + l_ids.append(sid) + sids[key_label] = l_ids + return sids + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + label = d[key] + if label.shape[0] != 1: + raise ValueError("Only supports single channel labels!") + + if len(label.shape) != 4: # only for 3D + raise ValueError("Only supports label with shape CHWD!") + + sids = self._apply(label, d) + if sids is not None and len(sids.keys()): + d[self.sids] = sids + return d + else: + print("This transform only applies to label key") + return d + + +class AddInitialSeedPointDeepEditd(Randomizable, MapTransform): + """ + Add random guidance as initial seed point for a given label. + + Note that the label is of size (C, D, H, W) or (C, H, W) + + The guidance is of size (2, N, # of dims) where N is number of guidance added. + # of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W) + + Args: + guidance: key to store guidance. + sids: key that represents lists of valid slice indices for the given label. + sid: key that represents the slice to add initial seed point. If not present, random sid will be chosen. + connected_regions: maximum connected regions to use for adding initial points. + """ + + def __init__( + self, + keys: KeysCollection, + guidance: str = "guidance", + sids: str = "sids", + sid: str = "sid", + connected_regions: int = 5, + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) + self.sids_key = sids + self.sid_key = sid + self.sid: Dict[str, int] = dict() + self.guidance = guidance + self.connected_regions = connected_regions + + def _apply(self, label, sid, key_label): + dimensions = 3 if len(label.shape) > 3 else 2 + self.default_guidance = [-1] * (dimensions + 1) + + dims = dimensions + if sid is not None and dimensions == 3: + dims = 2 + label = label[0][..., sid][np.newaxis] # Assume channel is first and depth is last CHWD + + # THERE MAY BE MULTIPLE BLOBS FOR SINGLE LABEL IN THE SELECTED SLICE + label = (label > 0.5).astype(np.float32) + # measure.label: Label connected regions of an integer array - Two pixels are connected + # when they are neighbors and have the same value + blobs_labels = measure.label(label.astype(int), background=0) if dims == 2 else label + if np.max(blobs_labels) <= 0: + raise AssertionError(f"SLICES NOT FOUND FOR LABEL: {key_label}") + + pos_guidance = [] + for ridx in range(1, 2 if dims == 3 else self.connected_regions + 1): + if dims == 2: + label = (blobs_labels == ridx).astype(np.float32) + if np.sum(label) == 0: + pos_guidance.append(self.default_guidance) + continue + + # The distance transform provides a metric or measure of the separation of points in the image. + # This function calculates the distance between each pixel that is set to off (0) and + # the nearest nonzero pixel for binary images - http://matlab.izmiran.ru/help/toolbox/images/morph14.html + distance = distance_transform_cdt(label).flatten() + probability = np.exp(distance) - 1.0 + + idx = np.where(label.flatten() > 0)[0] + seed = self.R.choice(idx, size=1, p=probability[idx] / np.sum(probability[idx])) + dst = distance[seed] - pos = np.multiply(pos_clicks, factor).astype(int, copy=False).tolist() if len(pos_clicks) else [] - neg = np.multiply(neg_clicks, factor).astype(int, copy=False).tolist() if len(neg_clicks) else [] + g = np.asarray(np.unravel_index(seed, label.shape)).transpose().tolist()[0] + g[0] = dst[0] # for debug + if dimensions == 2 or dims == 3: + pos_guidance.append(g) + else: + # Clicks are created using this convention Channel Height Width Depth (CHWD) + pos_guidance.append([g[0], g[-2], g[-1], sid]) # Assume channel is first and depth is last CHWD + + return np.asarray([pos_guidance]) + + def _randomize(self, d, key_label): + sids = d.get(self.sids_key).get(key_label) if d.get(self.sids_key) is not None else None + sid = d.get(self.sid_key).get(key_label) if d.get(self.sid_key) is not None else None + if sids is not None and sids: + if sid is None or sid not in sids: + sid = self.R.choice(sids, replace=False) + else: + logger.info(f"Not slice IDs for label: {key_label}") + sid = None + self.sid[key_label] = sid - d[self.guidance] = [pos, neg] + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + label_guidances = {} + for key_label in d["sids"].keys(): + # Randomize: Select a random slice + self._randomize(d, key_label) + # Generate guidance base on selected slice + tmp_label = np.copy(d[key]) + # Taking one label to create the guidance + if key_label != "background": + tmp_label[tmp_label != float(d["label_names"][key_label])] = 0 + else: + tmp_label[tmp_label != float(d["label_names"][key_label])] = 1 + tmp_label = 1 - tmp_label + label_guidances[key_label] = json.dumps( + self._apply(tmp_label, self.sid.get(key_label), key_label).astype(int).tolist() + ) + d[self.guidance] = label_guidances + return d + else: + print("This transform only applies to label key") return d -class ClickRatioAddRandomGuidanced(Randomizable, Transform): +class FindDiscrepancyRegionsDeepEditd(MapTransform): + """ + Find discrepancy between prediction and actual during click interactions during training. + + Args: + pred: key to prediction source. + discrepancy: key to store discrepancies found between label and prediction. + """ + + def __init__( + self, + keys: KeysCollection, + pred: str = "pred", + discrepancy: str = "discrepancy", + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) + self.pred = pred + self.discrepancy = discrepancy + + @staticmethod + def disparity(label, pred): + disparity = label - pred + # Negative ONES mean predicted label is not part of the ground truth + # Positive ONES mean predicted label missed that region of the ground truth + pos_disparity = (disparity > 0).astype(np.float32) + neg_disparity = (disparity < 0).astype(np.float32) + return [pos_disparity, neg_disparity] + + def _apply(self, label, pred): + return self.disparity(label, pred) + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + all_discrepancies = {} + for _, (key_label, val_label) in enumerate(d["label_names"].items()): + if key_label != "background": + # Taking single label + label = np.copy(d[key]) + label[label != val_label] = 0 + # Label should be represented in 1 + label = (label > 0.5).astype(np.float32) + # Taking single prediction + pred = np.copy(d[self.pred]) + pred[pred != val_label] = 0 + # Prediction should be represented in one + pred = (pred > 0.5).astype(np.float32) + else: + # Taking single label + label = np.copy(d[key]) + label[label != val_label] = 1 + label = 1 - label + # Label should be represented in 1 + label = (label > 0.5).astype(np.float32) + # Taking single prediction + pred = np.copy(d[self.pred]) + pred[pred != val_label] = 1 + pred = 1 - pred + # Prediction should be represented in one + pred = (pred > 0.5).astype(np.float32) + all_discrepancies[key_label] = self._apply(label, pred) + d[self.discrepancy] = all_discrepancies + return d + else: + print("This transform only applies to 'label' key") + return d + + +class AddRandomGuidanceDeepEditd(Randomizable, MapTransform): """ Add random guidance based on discrepancies that were found between label and prediction. + Args: guidance: key to guidance source, shape (2, N, # of dim) - discrepancy: key that represents discrepancies found between label and prediction, shape (2, C, D, H, W) or (2, C, H, W) - probability: key that represents click/interaction probability, shape (1) - fn_fp_click_ratio: ratio of clicks between FN and FP + discrepancy: key to discrepancy map between label and prediction shape (2, C, H, W, D) or (2, C, H, W) + probability: key to click/interaction probability, shape (1) """ def __init__( self, + keys: KeysCollection, guidance: str = "guidance", discrepancy: str = "discrepancy", probability: str = "probability", - fn_fp_click_ratio: Tuple[float, float] = (1.0, 1.0), + allow_missing_keys: bool = False, ): + super().__init__(keys, allow_missing_keys) self.guidance = guidance self.discrepancy = discrepancy self.probability = probability - self.fn_fp_click_ratio = fn_fp_click_ratio self._will_interact = None + self.is_pos = None + self.is_other = None + self.default_guidance = None def randomize(self, data=None): probability = data[self.probability] @@ -108,7 +506,7 @@ def randomize(self, data=None): def find_guidance(self, discrepancy): distance = distance_transform_cdt(discrepancy).flatten() - probability = np.exp(distance) - 1.0 + probability = np.exp(distance.flatten()) - 1.0 idx = np.where(discrepancy.flatten() > 0)[0] if np.sum(discrepancy > 0) > 0: @@ -120,51 +518,363 @@ def find_guidance(self, discrepancy): return g return None - def add_guidance(self, discrepancy, will_interact): - if not will_interact: - return None, None + def add_guidance(self, guidance, discrepancy, label_names, labels): - pos_discr = discrepancy[0] - neg_discr = discrepancy[1] + # Positive clicks of the segment in the iteration + pos_discr = discrepancy[0] # idx 0 is positive discrepancy and idx 1 is negative discrepancy - can_be_positive = np.sum(pos_discr) > 0 - can_be_negative = np.sum(neg_discr) > 0 + # Check the areas that belong to other segments + other_discrepancy_areas = {} + for _, (key_label, val_label) in enumerate(label_names.items()): + if key_label != "background": + tmp_label = np.copy(labels) + tmp_label[tmp_label != val_label] = 0 + tmp_label = (tmp_label > 0.5).astype(np.float32) + other_discrepancy_areas[key_label] = np.sum(discrepancy[1] * tmp_label) + else: + tmp_label = np.copy(labels) + tmp_label[tmp_label != val_label] = 1 + tmp_label = 1 - tmp_label + other_discrepancy_areas[key_label] = np.sum(discrepancy[1] * tmp_label) + + # Add guidance to the current key label + if np.sum(pos_discr) > 0: + guidance.append(self.find_guidance(pos_discr)) + self.is_pos = True + + # Add guidance to the other areas + for key_label in label_names.keys(): + # Areas that cover more than 50 voxels + if other_discrepancy_areas[key_label] > 50: + self.is_other = True + if key_label != "background": + tmp_label = np.copy(labels) + tmp_label[tmp_label != label_names[key_label]] = 0 + tmp_label = (tmp_label > 0.5).astype(np.float32) + self.tmp_guidance[key_label].append(self.find_guidance(discrepancy[1] * tmp_label)) + else: + tmp_label = np.copy(labels) + tmp_label[tmp_label != label_names[key_label]] = 1 + tmp_label = 1 - tmp_label + self.tmp_guidance[key_label].append(self.find_guidance(discrepancy[1] * tmp_label)) - pos_prob = self.fn_fp_click_ratio[0] / (self.fn_fp_click_ratio[0] + self.fn_fp_click_ratio[1]) - neg_prob = self.fn_fp_click_ratio[1] / (self.fn_fp_click_ratio[0] + self.fn_fp_click_ratio[1]) + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + guidance = d[self.guidance] + discrepancy = d[self.discrepancy] + self.randomize(data) + if self._will_interact: + # Convert all guidance to lists so new guidance can be easily appended + self.tmp_guidance = {} + for key_label in d["label_names"].keys(): + tmp_gui = guidance[key_label] + tmp_gui = tmp_gui.tolist() if isinstance(tmp_gui, np.ndarray) else tmp_gui + tmp_gui = json.loads(tmp_gui) if isinstance(tmp_gui, str) else tmp_gui + self.tmp_guidance[key_label] = [j for j in tmp_gui if -1 not in j] + + # Add guidance according to discrepancy + for key_label in d["label_names"].keys(): + # Add guidance based on discrepancy + self.add_guidance(self.tmp_guidance[key_label], discrepancy[key_label], d["label_names"], d["label"]) + + # Checking the number of clicks + num_clicks = random.randint(1, 10) + counter = 0 + keep_guidance = [] + while True: + aux_label = random.choice(list(d["label_names"].keys())) + if aux_label in keep_guidance: + pass + else: + keep_guidance.append(aux_label) + counter = counter + len(self.tmp_guidance[aux_label]) + # If collected clicks is bigger than max clicks, discard the others + if counter >= num_clicks: + for key_label in d["label_names"].keys(): + if key_label not in keep_guidance: + self.tmp_guidance[key_label] = [] + logger.info(f"Number of simulated clicks: {counter}") + break + + # Breaking once all labels are covered + if len(keep_guidance) == len(d["label_names"].keys()): + logger.info(f"Number of simulated clicks: {counter}") + break - correct_pos = self.R.choice([True, False], p=[pos_prob, neg_prob]) + return d - if can_be_positive and not can_be_negative: - return self.find_guidance(pos_discr), None - if not can_be_positive and can_be_negative: - return None, self.find_guidance(neg_discr) +class AddGuidanceFromPointsDeepEditd(Transform): + """ + Add guidance based on user clicks. ONLY WORKS FOR 3D - if correct_pos and can_be_positive: - return self.find_guidance(pos_discr), None + We assume the input is loaded by LoadImaged and has the shape of (H, W, D) originally. + Clicks always specify the coordinates in (H, W, D) - if not correct_pos and can_be_negative: - return None, self.find_guidance(neg_discr) - return None, None + Args: + ref_image: key to reference image to fetch current and original image details. + guidance: output key to store guidance. + meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`. + for example, for data with key `image`, the metadata by default is in `image_meta_dict`. + the metadata is a dictionary object which contains: filename, original_shape, etc. + if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`. + meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. + For example, to handle key `image`, read/write affine matrices from the + metadata `image_meta_dict` dictionary's `affine` field. - def _apply(self, guidance, discrepancy): - guidance = guidance.tolist() if isinstance(guidance, np.ndarray) else guidance - guidance = json.loads(guidance) if isinstance(guidance, str) else guidance - pos, neg = self.add_guidance(discrepancy, self._will_interact) - if pos: - guidance[0].append(pos) - guidance[1].append([-1] * len(pos)) - if neg: - guidance[0].append([-1] * len(neg)) - guidance[1].append(neg) + """ - return json.dumps(np.asarray(guidance, dtype=int).tolist()) + def __init__( + self, + ref_image, + guidance: str = "guidance", + label_names=None, + meta_keys: Optional[str] = None, + meta_key_postfix: str = "meta_dict", + ): + self.ref_image = ref_image + self.guidance = guidance + self.label_names = label_names + self.meta_keys = meta_keys + self.meta_key_postfix = meta_key_postfix + + @staticmethod + def _apply(clicks, factor): + if len(clicks): + guidance = np.multiply(clicks, factor).astype(int).tolist() + return guidance + else: + return [] def __call__(self, data): d = dict(data) - guidance = d[self.guidance] - discrepancy = d[self.discrepancy] - self.randomize(data) - d[self.guidance] = self._apply(guidance, discrepancy) + meta_dict_key = self.meta_keys or f"{self.ref_image}_{self.meta_key_postfix}" + if meta_dict_key not in d: + raise RuntimeError(f"Missing meta_dict {meta_dict_key} in data!") + if "spatial_shape" not in d[meta_dict_key]: + raise RuntimeError('Missing "spatial_shape" in meta_dict!') + + # Assume channel is first and depth is last CHWD + original_shape = d[meta_dict_key]["spatial_shape"] + current_shape = list(d[self.ref_image].shape)[1:] + + # in here we assume the depth dimension is in the last dimension of "original_shape" and "current_shape" + factor = np.array(current_shape) / original_shape + + # Creating guidance for all clicks + all_guidances = {} + for key_label in self.label_names.keys(): + clicks = d.get(key_label, []) + clicks = list(np.array(clicks).astype(int)) + all_guidances[key_label] = self._apply(clicks, factor) + d[self.guidance] = all_guidances + return d + + +class ResizeGuidanceMultipleLabelDeepEditd(Transform): + """ + Resize the guidance based on cropped vs resized image. + + """ + + def __init__(self, guidance: str, ref_image: str) -> None: + self.guidance = guidance + self.ref_image = ref_image + + def __call__(self, data): + d = dict(data) + # Assume channel is first and depth is last CHWD + current_shape = d[self.ref_image].shape[1:] + original_shape = d["image_meta_dict"]["spatial_shape"] + + factor = np.divide(current_shape, original_shape) + all_guidances = {} + for key_label in d[self.guidance].keys(): + guidance = ( + np.multiply(d[self.guidance][key_label], factor).astype(int).tolist() + if len(d[self.guidance][key_label]) + else [] + ) + all_guidances[key_label] = guidance + + d[self.guidance] = all_guidances + return d + + +class SplitPredsLabeld(MapTransform): + """ + Split preds and labels for individual evaluation + + """ + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "pred": + for idx, (key_label, _) in enumerate(d["label_names"].items()): + if key_label != "background": + d[f"pred_{key_label}"] = d[key][idx + 1, ...][None] + d[f"label_{key_label}"] = d["label"][idx + 1, ...][None] + elif key != "pred": + logger.info("This is only for pred key") + return d + + +class AddInitialSeedPointMissingLabelsd(Randomizable, MapTransform): + """ + Add random guidance as initial seed point for a given label. + Note that the label is of size (C, D, H, W) or (C, H, W) + The guidance is of size (2, N, # of dims) where N is number of guidance added. + # of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W) + Args: + guidance: key to store guidance. + sids: key that represents lists of valid slice indices for the given label. + sid: key that represents the slice to add initial seed point. If not present, random sid will be chosen. + connected_regions: maximum connected regions to use for adding initial points. + """ + + def __init__( + self, + keys: KeysCollection, + guidance: str = "guidance", + sids: str = "sids", + sid: str = "sid", + connected_regions: int = 5, + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) + self.sids_key = sids + self.sid_key = sid + self.sid: Dict[str, int] = dict() + self.guidance = guidance + self.connected_regions = connected_regions + + def _apply(self, label, sid): + dimensions = 3 if len(label.shape) > 3 else 2 + self.default_guidance = [-1] * (dimensions + 1) + + dims = dimensions + if sid is not None and dimensions == 3: + dims = 2 + label = label[0][..., sid][np.newaxis] # Assume channel is first and depth is last CHWD + + # THERE MAY BE MULTIPLE BLOBS FOR SINGLE LABEL IN THE SELECTED SLICE + label = (label > 0.5).astype(np.float32) + # measure.label: Label connected regions of an integer array - Two pixels are connected + # when they are neighbors and have the same value + blobs_labels = measure.label(label.astype(int), background=0) if dims == 2 else label + + label_guidance = [] + # If there are is presence of that label in this slice + if np.max(blobs_labels) <= 0: + label_guidance.append(self.default_guidance) + else: + for ridx in range(1, 2 if dims == 3 else self.connected_regions + 1): + if dims == 2: + label = (blobs_labels == ridx).astype(np.float32) + if np.sum(label) == 0: + label_guidance.append(self.default_guidance) + continue + + # The distance transform provides a metric or measure of the separation of points in the image. + # This function calculates the distance between each pixel that is set to off (0) and + # the nearest nonzero pixel for binary images + # http://matlab.izmiran.ru/help/toolbox/images/morph14.html + distance = distance_transform_cdt(label).flatten() + probability = np.exp(distance) - 1.0 + + idx = np.where(label.flatten() > 0)[0] + seed = self.R.choice(idx, size=1, p=probability[idx] / np.sum(probability[idx])) + dst = distance[seed] + + g = np.asarray(np.unravel_index(seed, label.shape)).transpose().tolist()[0] + g[0] = dst[0] # for debug + if dimensions == 2 or dims == 3: + label_guidance.append(g) + else: + # Clicks are created using this convention Channel Height Width Depth (CHWD) + label_guidance.append([g[0], g[-2], g[-1], sid]) # Assume channel is first and depth is last CHWD + + return np.asarray(label_guidance) + + def _randomize(self, d, key_label): + sids = d.get(self.sids_key).get(key_label) if d.get(self.sids_key) is not None else None + sid = d.get(self.sid_key).get(key_label) if d.get(self.sid_key) is not None else None + if sids is not None and sids: + if sid is None or sid not in sids: + sid = self.R.choice(sids, replace=False) + else: + logger.info(f"Not slice IDs for label: {key_label}") + sid = None + self.sid[key_label] = sid + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + label_guidances = {} + for key_label in d["sids"].keys(): + # Randomize: Select a random slice + self._randomize(d, key_label) + # Generate guidance base on selected slice + tmp_label = np.copy(d[key]) + # Taking one label to create the guidance + if key_label != "background": + tmp_label[tmp_label != float(d["label_names"][key_label])] = 0 + else: + tmp_label[tmp_label != float(d["label_names"][key_label])] = 1 + tmp_label = 1 - tmp_label + label_guidances[key_label] = json.dumps( + self._apply(tmp_label, self.sid.get(key_label)).astype(int).tolist() + ) + d[self.guidance] = label_guidances + return d + else: + print("This transform only applies to label key") + return d + + +class FindAllValidSlicesMissingLabelsd(MapTransform): + """ + Find/List all valid slices in the labels. + Label is assumed to be a 4D Volume with shape CHWD, where C=1. + Args: + sids: key to store slices indices having valid label map. + """ + + def __init__(self, keys: KeysCollection, sids="sids", allow_missing_keys: bool = False): + super().__init__(keys, allow_missing_keys) + self.sids = sids + + def _apply(self, label, d): + sids = {} + for key_label in d["label_names"].keys(): + l_ids = [] + for sid in range(label.shape[-1]): # Assume channel is first and depth is last CHWD + if d["label_names"][key_label] in label[0][..., sid]: + l_ids.append(sid) + # If there are not slices with the label + if l_ids == []: + l_ids = [-1] * 10 + sids[key_label] = l_ids + return sids + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d: Dict = dict(data) + for key in self.key_iterator(d): + if key == "label": + label = d[key] + if label.shape[0] != 1: + raise ValueError("Only supports single channel labels!") + + if len(label.shape) != 4: # only for 3D + raise ValueError("Only supports label with shape CHWD!") + + sids = self._apply(label, d) + if sids is not None and len(sids.keys()): + d[self.sids] = sids + return d + else: + print("This transform only applies to label key") return d diff --git a/monai/apps/deepgrow/dataset.py b/monai/apps/deepgrow/dataset.py index 721781196b..d51fd4e238 100644 --- a/monai/apps/deepgrow/dataset.py +++ b/monai/apps/deepgrow/dataset.py @@ -15,8 +15,9 @@ import numpy as np -from monai.transforms import AsChannelFirstd, Compose, LoadImaged, Orientationd, Spacingd +from monai.transforms import AsChannelFirstd, Compose, FromMetaTensord, LoadImaged, Orientationd, Spacingd, ToNumpyd from monai.utils import GridSampleMode +from monai.utils.enums import PostFix def create_dataset( @@ -128,6 +129,8 @@ def _default_transforms(image_key, label_key, pixdim): AsChannelFirstd(keys=keys), Orientationd(keys=keys, axcodes="RAS"), Spacingd(keys=keys, pixdim=pixdim, mode=mode), + FromMetaTensord(keys=keys), + ToNumpyd(keys=keys + [PostFix.meta(k) for k in keys]), ] ) diff --git a/monai/apps/deepgrow/transforms.py b/monai/apps/deepgrow/transforms.py index 6da614f46c..c439469aea 100644 --- a/monai/apps/deepgrow/transforms.py +++ b/monai/apps/deepgrow/transforms.py @@ -372,13 +372,13 @@ class SpatialCropForegroundd(MapTransform): allow_smaller: when computing box size with `margin`, whether allow the image size to be smaller than box size, default to `True`. if the margined size is bigger than image size, will pad with specified `mode`. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `{key}_{meta_key_postfix}` to fetch/store the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_keys is None, use `{key}_{meta_key_postfix}` to fetch/store the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. start_coord_key: key to record the start coordinate of spatial bounding box for foreground. @@ -475,12 +475,12 @@ class AddGuidanceFromPointsd(Transform): depth_first: if depth (slices) is positioned at first dimension. spatial_dims: dimensions based on model used for deepgrow (2D vs 3D). slice_key: key that represents applicable slice to add guidance. - meta_keys: explicitly indicate the key of the meta data dictionary of `ref_image`. + meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`. - meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. @@ -589,13 +589,13 @@ class SpatialCropGuidanced(MapTransform): guidance: key to the guidance. It is used to generate the bounding box of foreground spatial_size: minimal spatial size of the image patch e.g. [128, 128, 128] to fit in. margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. start_coord_key: key to record the start coordinate of spatial bounding box for foreground. @@ -712,12 +712,12 @@ class ResizeGuidanced(Transform): Args: guidance: key to guidance ref_image: key to reference image to fetch current and original image details - meta_keys: explicitly indicate the key of the meta data dictionary of `ref_image`. + meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`. - meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. cropped_shape_key: key that records cropped shape for foreground. @@ -787,13 +787,13 @@ class RestoreLabeld(MapTransform): align_corners: Geometrically, we consider the pixels of the input as squares rather than points. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html It also can be a sequence of bool, each element corresponds to a key in ``keys``. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_key is None, use `key_{meta_key_postfix} to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_key is None, use `key_{meta_key_postfix} to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. start_coord_key: key that records the start coordinate of spatial bounding box for foreground. @@ -897,13 +897,13 @@ class Fetch2DSliced(MapTransform): keys: keys of the corresponding items to be transformed. guidance: key that represents guidance. axis: axis that represents slice in 3D volume. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: use `key_{meta_key_postfix}` to fetch the meta data according to the key data, - default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: use `key_{meta_key_postfix}` to fetch the metadata according to the key data, + default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. allow_missing_keys: don't raise exception if key is missing. diff --git a/monai/apps/detection/__init__.py b/monai/apps/detection/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/__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/monai/apps/detection/metrics/__init__.py b/monai/apps/detection/metrics/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/metrics/__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/monai/apps/detection/metrics/coco.py b/monai/apps/detection/metrics/coco.py new file mode 100644 index 0000000000..7c6b02fe4f --- /dev/null +++ b/monai/apps/detection/metrics/coco.py @@ -0,0 +1,547 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/evaluator/detection/coco.py +# which has the following license... +# https://github.com/MIC-DKFZ/nnDetection/blob/main/LICENSE +# +# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany +# 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. + +# ========================================================================= +# Adapted from https://github.com/cocodataset/cocoapi +# which has the following license... +# https://github.com/cocodataset/cocoapi/blob/master/license.txt + +# Copyright (c) 2014, Piotr Dollar and Tsung-Yi Lin +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# The views and conclusions contained in the software and documentation are those +# of the authors and should not be interpreted as representing official policies, +# either expressed or implied, of the FreeBSD Project. + +""" +This script is almost same with https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/evaluator/detection/coco.py +The changes include 1) code reformatting, 2) docstrings. +""" + +import logging as logger +import time +from typing import Dict, List, Sequence, Tuple, Union + +import numpy as np + + +class COCOMetric: + def __init__( + self, + classes: Sequence[str], + iou_list: Sequence[float] = (0.1, 0.5, 0.75), + iou_range: Sequence[float] = (0.1, 0.5, 0.05), + max_detection: Sequence[int] = (1, 5, 100), + per_class: bool = True, + verbose: bool = True, + ): + """ + Class to compute COCO metrics + Metrics computed includes, + + - mAP over the IoU range specified by `iou_range` at last value of `max_detection` + - AP values at IoU thresholds specified by `iou_list` at last value of `max_detection` + - AR over max detections thresholds defined by `max_detection` (over iou range) + + Args: + classes (Sequence[str]): name of each class (index needs to correspond to predicted class indices!) + iou_list (Sequence[float]): specific thresholds where ap is evaluated and saved + iou_range (Sequence[float]): (start, stop, step) for mAP iou thresholds + max_detection (Sequence[int]): maximum number of detections per image + verbose (bool): log time needed for evaluation + + Example: + + .. code-block:: python + + from monai.data.box_utils import box_iou + from monai.apps.detection.metrics.coco import COCOMetric + from monai.apps.detection.metrics.matching import matching_batch + # 3D example outputs of one image from detector + val_outputs_all = [ + {"boxes": torch.tensor([[1,1,1,3,4,5]],dtype=torch.float16), + "labels": torch.randint(3,(1,)), + "scores": torch.randn((1,)).absolute()}, + ] + val_targets_all = [ + {"boxes": torch.tensor([[1,1,1,2,6,4]],dtype=torch.float16), + "labels": torch.randint(3,(1,))}, + ] + + coco_metric = COCOMetric( + classes=['c0','c1','c2'], iou_list=[0.1], max_detection=[10] + ) + results_metric = matching_batch( + iou_fn=box_iou, + iou_thresholds=coco_metric.iou_thresholds, + pred_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_outputs_all], + pred_classes=[val_data_i["labels"].numpy() for val_data_i in val_outputs_all], + pred_scores=[val_data_i["scores"].numpy() for val_data_i in val_outputs_all], + gt_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_targets_all], + gt_classes=[val_data_i["labels"].numpy() for val_data_i in val_targets_all], + ) + val_metric_dict = coco_metric(results_metric) + print(val_metric_dict) + """ + self.verbose = verbose + self.classes = classes + self.per_class = per_class + + iou_list_np = np.array(iou_list) + _iou_range = np.linspace( + iou_range[0], iou_range[1], int(np.round((iou_range[1] - iou_range[0]) / iou_range[2])) + 1, endpoint=True + ) + self.iou_thresholds = np.union1d(iou_list_np, _iou_range) + self.iou_range = iou_range + + # get indices of iou values of ious range and ious list for later evaluation + self.iou_list_idx = np.nonzero(iou_list_np[:, np.newaxis] == self.iou_thresholds[np.newaxis])[1] + self.iou_range_idx = np.nonzero(_iou_range[:, np.newaxis] == self.iou_thresholds[np.newaxis])[1] + + if ( + not (self.iou_thresholds[self.iou_list_idx] == iou_list_np).all() + or not (self.iou_thresholds[self.iou_range_idx] == _iou_range).all() + ): + raise ValueError( + "Require self.iou_thresholds[self.iou_list_idx] == iou_list_np and " + "self.iou_thresholds[self.iou_range_idx] == _iou_range." + ) + + self.recall_thresholds = np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True) + self.max_detections = max_detection + + def __call__(self, *args, **kwargs) -> Tuple[Dict[str, float], Union[Dict[str, np.ndarray], None]]: + """ + Compute metric. See :func:`compute` for more information. + + Args: + *args: positional arguments passed to :func:`compute` + **kwargs: keyword arguments passed to :func:`compute` + + Returns: + Dict[str, float]: dictionary with scalar values for evaluation + Dict[str, np.ndarray]: dictionary with arrays, e.g. for visualization of graphs + """ + return self.compute(*args, **kwargs) + + def check_number_of_iou(self, *args) -> None: + """ + Check if shape of input in first dimension is consistent with expected IoU values + (assumes IoU dimension is the first dimension) + + Args: + args: array like inputs with shape function + """ + num_ious = len(self.get_iou_thresholds()) + for arg in args: + if arg.shape[0] != num_ious: + raise ValueError( + f"Require arg.shape[0] == len(self.get_iou_thresholds()). Got arg.shape[0]={arg.shape[0]}, " + f"self.get_iou_thresholds()={self.get_iou_thresholds()}." + ) + + def get_iou_thresholds(self) -> Sequence[float]: + """ + Return IoU thresholds needed for this metric in an numpy array + + Returns: + Sequence[float]: IoU thresholds [M], M is the number of thresholds + """ + return list(self.iou_thresholds) + + def compute(self, results_list: List[Dict[int, Dict[str, np.ndarray]]]) -> Tuple[Dict[str, float], None]: + """ + Compute COCO metrics + + Args: + results_list (List[Dict[int, Dict[str, np.ndarray]]]): list with results per image (in list) + per category (dict). Inner Dict contains multiple results obtained by :func:`box_matching_batch`. + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + + Returns: + Dict[str, float], dictionary with coco metrics + """ + if self.verbose: + logger.info("Start COCO metric computation...") + tic = time.time() + + dataset_statistics = self._compute_statistics(results_list=results_list) # Dict[str, Union[np.ndarray, List]] + + if self.verbose: + toc = time.time() + logger.info(f"Statistics for COCO metrics finished (t={(toc - tic):0.2f}s).") + + results = {} + results.update(self._compute_ap(dataset_statistics)) + results.update(self._compute_ar(dataset_statistics)) + + if self.verbose: + toc = time.time() + logger.info(f"COCO metrics computed in t={(toc - tic):0.2f}s.") + return results, None + + def _compute_ap(self, dataset_statistics: Dict[str, Union[np.ndarray, List]]) -> Dict[str, float]: + """ + Compute AP metrics + + Args: + dataset_statistics (List[Dict[int, Dict[str, np.ndarray]]]): list with result s per image (in list) + per category (dict). Inner Dict contains multiple results obtained by :func:`box_matching_batch`. + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + """ + results = {} + if self.iou_range: # mAP + key = ( + f"mAP_IoU_{self.iou_range[0]:.2f}_{self.iou_range[1]:.2f}_{self.iou_range[2]:.2f}_" + f"MaxDet_{self.max_detections[-1]}" + ) + results[key] = self._select_ap(dataset_statistics, iou_idx=self.iou_range_idx, max_det_idx=-1) + + if self.per_class: + for cls_idx, cls_str in enumerate(self.classes): # per class results + key = ( + f"{cls_str}_" + f"mAP_IoU_{self.iou_range[0]:.2f}_{self.iou_range[1]:.2f}_{self.iou_range[2]:.2f}_" + f"MaxDet_{self.max_detections[-1]}" + ) + results[key] = self._select_ap( + dataset_statistics, iou_idx=self.iou_range_idx, cls_idx=cls_idx, max_det_idx=-1 + ) + + for idx in self.iou_list_idx: # AP@IoU + key = f"AP_IoU_{self.iou_thresholds[idx]:.2f}_MaxDet_{self.max_detections[-1]}" + results[key] = self._select_ap(dataset_statistics, iou_idx=[idx], max_det_idx=-1) + + if self.per_class: + for cls_idx, cls_str in enumerate(self.classes): # per class results + key = f"{cls_str}_" f"AP_IoU_{self.iou_thresholds[idx]:.2f}_" f"MaxDet_{self.max_detections[-1]}" + results[key] = self._select_ap(dataset_statistics, iou_idx=[idx], cls_idx=cls_idx, max_det_idx=-1) + return results + + def _compute_ar(self, dataset_statistics: Dict[str, Union[np.ndarray, List]]) -> Dict[str, float]: + """ + Compute AR metrics + + Args: + dataset_statistics (List[Dict[int, Dict[str, np.ndarray]]]): list with result s per image (in list) + per category (dict). Inner Dict contains multiple results obtained by :func:`box_matching_batch`. + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + """ + results = {} + for max_det_idx, max_det in enumerate(self.max_detections): # mAR + key = f"mAR_IoU_{self.iou_range[0]:.2f}_{self.iou_range[1]:.2f}_{self.iou_range[2]:.2f}_MaxDet_{max_det}" + results[key] = self._select_ar(dataset_statistics, max_det_idx=max_det_idx) + + if self.per_class: + for cls_idx, cls_str in enumerate(self.classes): # per class results + key = ( + f"{cls_str}_" + f"mAR_IoU_{self.iou_range[0]:.2f}_{self.iou_range[1]:.2f}_{self.iou_range[2]:.2f}_" + f"MaxDet_{max_det}" + ) + results[key] = self._select_ar(dataset_statistics, cls_idx=cls_idx, max_det_idx=max_det_idx) + + for idx in self.iou_list_idx: # AR@IoU + key = f"AR_IoU_{self.iou_thresholds[idx]:.2f}_MaxDet_{self.max_detections[-1]}" + results[key] = self._select_ar(dataset_statistics, iou_idx=idx, max_det_idx=-1) + + if self.per_class: + for cls_idx, cls_str in enumerate(self.classes): # per class results + key = f"{cls_str}_" f"AR_IoU_{self.iou_thresholds[idx]:.2f}_" f"MaxDet_{self.max_detections[-1]}" + results[key] = self._select_ar(dataset_statistics, iou_idx=idx, cls_idx=cls_idx, max_det_idx=-1) + return results + + @staticmethod + def _select_ap( + dataset_statistics: dict, + iou_idx: Union[int, List[int], np.ndarray, None] = None, + cls_idx: Union[int, Sequence[int], None] = None, + max_det_idx: int = -1, + ) -> float: + """ + Compute average precision + + Args: + dataset_statistics (dict): computed statistics over dataset + + - `counts`: Number of thresholds, Number recall thresholds, Number of classes, Number of max + detection thresholds + - `recall`: Computed recall values [num_iou_th, num_classes, num_max_detections] + - `precision`: Precision values at specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + - `scores`: Scores corresponding to specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + iou_idx: index of IoU values to select for evaluation(if None, all values are used) + cls_idx: class indices to select, if None all classes will be selected + max_det_idx (int): index to select max detection threshold from data + + Returns: + np.ndarray: AP value + """ + prec = dataset_statistics["precision"] + if iou_idx is not None: + prec = prec[iou_idx] + if cls_idx is not None: + prec = prec[..., cls_idx, :] + prec = prec[..., max_det_idx] + return float(np.mean(prec)) + + @staticmethod + def _select_ar( + dataset_statistics: dict, + iou_idx: Union[int, Sequence[int], None] = None, + cls_idx: Union[int, Sequence[int], None] = None, + max_det_idx: int = -1, + ) -> float: + """ + Compute average recall + + Args: + dataset_statistics (dict): computed statistics over dataset + + - `counts`: Number of thresholds, Number recall thresholds, Number of classes, Number of max + detection thresholds + - `recall`: Computed recall values [num_iou_th, num_classes, num_max_detections] + - `precision`: Precision values at specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + - `scores`: Scores corresponding to specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + iou_idx: index of IoU values to select for evaluation(if None, all values are used) + cls_idx: class indices to select, if None all classes will be selected + max_det_idx (int): index to select max detection threshold from data + + Returns: + np.ndarray: recall value + """ + rec = dataset_statistics["recall"] + if iou_idx is not None: + rec = rec[iou_idx] + if cls_idx is not None: + rec = rec[..., cls_idx, :] + rec = rec[..., max_det_idx] + + if len(rec[rec > -1]) == 0: + return -1.0 + + return float(np.mean(rec[rec > -1])) + + def _compute_statistics( + self, results_list: List[Dict[int, Dict[str, np.ndarray]]] + ) -> Dict[str, Union[np.ndarray, List]]: + """ + Compute statistics needed for COCO metrics (mAP, AP of individual classes, mAP@IoU_Thresholds, AR) + Adapted from https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py + + Args: + results_list (List[Dict[int, Dict[str, np.ndarray]]]): list with result s per image (in list) + per cateory (dict). Inner Dict contains multiple results obtained by :func:`box_matching_batch`. + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + + Returns: + dict: computed statistics over dataset + - `counts`: Number of thresholds, Number recall thresholds, Number of classes, Number of max + detection thresholds + - `recall`: Computed recall values [num_iou_th, num_classes, num_max_detections] + - `precision`: Precision values at specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + - `scores`: Scores corresponding to specified recall thresholds + [num_iou_th, num_recall_th, num_classes, num_max_detections] + """ + num_iou_th = len(self.iou_thresholds) + num_recall_th = len(self.recall_thresholds) + num_classes = len(self.classes) + num_max_detections = len(self.max_detections) + + # -1 for the precision of absent categories + precision = -np.ones((num_iou_th, num_recall_th, num_classes, num_max_detections)) + recall = -np.ones((num_iou_th, num_classes, num_max_detections)) + scores = -np.ones((num_iou_th, num_recall_th, num_classes, num_max_detections)) + + for cls_idx, cls_i in enumerate(self.classes): # for each class + for max_det_idx, max_det in enumerate(self.max_detections): # for each maximum number of detections + results = [r[cls_idx] for r in results_list if cls_idx in r] # len is num_images + + if len(results) == 0: + logger.warning(f"WARNING, no results found for coco metric for class {cls_i}") + continue + + dt_scores = np.concatenate([r["dtScores"][0:max_det] for r in results]) + # different sorting method generates slightly different results. + # mergesort is used to be consistent as Matlab implementation. + inds = np.argsort(-dt_scores, kind="mergesort") + dt_scores_sorted = dt_scores[inds] + + # r['dtMatches'] [T, R], where R = sum(all detections) + dt_matches = np.concatenate([r["dtMatches"][:, 0:max_det] for r in results], axis=1)[:, inds] + dt_ignores = np.concatenate([r["dtIgnore"][:, 0:max_det] for r in results], axis=1)[:, inds] + self.check_number_of_iou(dt_matches, dt_ignores) + gt_ignore = np.concatenate([r["gtIgnore"] for r in results]) + num_gt = np.count_nonzero(gt_ignore == 0) # number of ground truth boxes (non ignored) + if num_gt == 0: + logger.warning(f"WARNING, no gt found for coco metric for class {cls_i}") + continue + + # ignore cases need to be handled differently for tp and fp + tps = np.logical_and(dt_matches, np.logical_not(dt_ignores)) + fps = np.logical_and(np.logical_not(dt_matches), np.logical_not(dt_ignores)) + + tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float32) + fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float32) + + for th_ind, (tp, fp) in enumerate(zip(tp_sum, fp_sum)): # for each threshold th_ind + tp, fp = np.array(tp), np.array(fp) + r, p, s = _compute_stats_single_threshold(tp, fp, dt_scores_sorted, self.recall_thresholds, num_gt) + recall[th_ind, cls_idx, max_det_idx] = r + precision[th_ind, :, cls_idx, max_det_idx] = p + # corresponding score thresholds for recall steps + scores[th_ind, :, cls_idx, max_det_idx] = s + + return { + "counts": [num_iou_th, num_recall_th, num_classes, num_max_detections], # [4] + "recall": recall, # [num_iou_th, num_classes, num_max_detections] + "precision": precision, # [num_iou_th, num_recall_th, num_classes, num_max_detections] + "scores": scores, # [num_iou_th, num_recall_th, num_classes, num_max_detections] + } + + +def _compute_stats_single_threshold( + tp: np.ndarray, + fp: np.ndarray, + dt_scores_sorted: np.ndarray, + recall_thresholds: Union[np.ndarray, Sequence[float]], + num_gt: int, +) -> Tuple[float, np.ndarray, np.ndarray]: + """ + Compute recall value, precision curve and scores thresholds + Adapted from https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py + + Args: + tp (np.ndarray): cumsum over true positives [R], R is the number of detections + fp (np.ndarray): cumsum over false positives [R], R is the number of detections + dt_scores_sorted (np.ndarray): sorted (descending) scores [R], R is the number of detections + recall_thresholds (Sequence[float]): recall thresholds which should be evaluated + num_gt (int): number of ground truth bounding boxes (excluding boxes which are ignored) + + Returns: + - float, overall recall for given IoU value + - np.ndarray, precision values at defined recall values + [RTH], where RTH is the number of recall thresholds + - np.ndarray, prediction scores corresponding to recall values + [RTH], where RTH is the number of recall thresholds + """ + num_recall_th = len(recall_thresholds) + + rc = tp / num_gt + # np.spacing(1) is the smallest representable epsilon with float + pr = tp / (fp + tp + np.spacing(1)) + + if len(tp): + recall = rc[-1] + else: + # no prediction + recall = 0 + + # array where precision values nearest to given recall th are saved + precision = np.zeros((num_recall_th,)) + # save scores for corresponding recall value in here + th_scores = np.zeros((num_recall_th,)) + # numpy is slow without cython optimization for accessing elements + # use python array gets significant speed improvement + pr = pr.tolist() + precision = precision.tolist() + + # smooth precision curve (create box shape) + for i in range(len(tp) - 1, 0, -1): + if pr[i] > pr[i - 1]: + pr[i - 1] = pr[i] + + # get indices to nearest given recall threshold (nn interpolation!) + inds = np.searchsorted(rc, recall_thresholds, side="left") + try: + for save_idx, array_index in enumerate(inds): + precision[save_idx] = pr[array_index] + th_scores[save_idx] = dt_scores_sorted[array_index] + except BaseException: + pass + + return recall, np.array(precision), np.array(th_scores) diff --git a/monai/apps/detection/metrics/matching.py b/monai/apps/detection/metrics/matching.py new file mode 100644 index 0000000000..6df026bf54 --- /dev/null +++ b/monai/apps/detection/metrics/matching.py @@ -0,0 +1,367 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/evaluator/detection/matching.py +# which has the following license... +# https://github.com/MIC-DKFZ/nnDetection/blob/main/LICENSE +# +# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany +# 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. + +# ========================================================================= +# Adapted from https://github.com/cocodataset/cocoapi +# which has the following license... +# https://github.com/cocodataset/cocoapi/blob/master/license.txt + +# Copyright (c) 2014, Piotr Dollar and Tsung-Yi Lin +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# The views and conclusions contained in the software and documentation are those +# of the authors and should not be interpreted as representing official policies, +# either expressed or implied, of the FreeBSD Project. + +""" +This script is almost same with https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/evaluator/detection/matching.py +The changes include 1) code reformatting, 2) docstrings, +3) allow input args gt_ignore to be optional. (If so, no GT boxes will be ignored.) +""" + +from typing import Callable, Dict, List, Sequence, Union + +import numpy as np + +__all__ = ["matching_batch"] + + +def matching_batch( + iou_fn: Callable[[np.ndarray, np.ndarray], np.ndarray], + iou_thresholds: Sequence[float], + pred_boxes: Sequence[np.ndarray], + pred_classes: Sequence[np.ndarray], + pred_scores: Sequence[np.ndarray], + gt_boxes: Sequence[np.ndarray], + gt_classes: Sequence[np.ndarray], + gt_ignore: Union[Sequence[Sequence[bool]], Sequence[np.ndarray], None] = None, + max_detections: int = 100, +) -> List[Dict[int, Dict[str, np.ndarray]]]: + """ + Match boxes of a batch to corresponding ground truth for each category + independently. + + Args: + iou_fn: compute overlap for each pair + iou_thresholds: defined which IoU thresholds should be evaluated + pred_boxes: predicted boxes from single batch; List[[D, dim * 2]], + D number of predictions + pred_classes: predicted classes from a single batch; List[[D]], + D number of predictions + pred_scores: predicted score for each bounding box; List[[D]], + D number of predictions + gt_boxes: ground truth boxes; List[[G, dim * 2]], G number of ground + truth + gt_classes: ground truth classes; List[[G]], G number of ground truth + gt_ignore: specified if which ground truth boxes are not counted as + true positives. If not given, when use all the gt_boxes. + (detections which match theses boxes are not counted as false + positives either); List[[G]], G number of ground truth + max_detections: maximum number of detections which should be evaluated + + Returns: + List[Dict[int, Dict[str, np.ndarray]]], each Dict[str, np.ndarray] corresponds to an image. + Dict has the following keys. + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + + Example: + + .. code-block:: python + + from monai.data.box_utils import box_iou + from monai.apps.detection.metrics.coco import COCOMetric + from monai.apps.detection.metrics.matching import matching_batch + # 3D example outputs of one image from detector + val_outputs_all = [ + {"boxes": torch.tensor([[1,1,1,3,4,5]],dtype=torch.float16), + "labels": torch.randint(3,(1,)), + "scores": torch.randn((1,)).absolute()}, + ] + val_targets_all = [ + {"boxes": torch.tensor([[1,1,1,2,6,4]],dtype=torch.float16), + "labels": torch.randint(3,(1,))}, + ] + + coco_metric = COCOMetric( + classes=['c0','c1','c2'], iou_list=[0.1], max_detection=[10] + ) + results_metric = matching_batch( + iou_fn=box_iou, + iou_thresholds=coco_metric.iou_thresholds, + pred_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_outputs_all], + pred_classes=[val_data_i["labels"].numpy() for val_data_i in val_outputs_all], + pred_scores=[val_data_i["scores"].numpy() for val_data_i in val_outputs_all], + gt_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_targets_all], + gt_classes=[val_data_i["labels"].numpy() for val_data_i in val_targets_all], + ) + val_metric_dict = coco_metric(results_metric) + print(val_metric_dict) + """ + results = [] + if gt_ignore is None: + gt_ignore = [np.full_like(gt_c, False) for gt_c in gt_classes] + # iterate over images/batches + for pboxes, pclasses, pscores, gboxes, gclasses, gignore in zip( + pred_boxes, pred_classes, pred_scores, gt_boxes, gt_classes, gt_ignore + ): + # for each image + img_classes = np.union1d(pclasses, gclasses) # possible class labels + result = {} # dict contains results for each class in one image + for c in img_classes: + pred_mask = pclasses == c # bool mask predictions with current class + gt_mask = gclasses == c # nool mask ground trtuh with current class + + if not np.any(gt_mask): # no ground truth + result[c] = _matching_no_gt( + iou_thresholds=iou_thresholds, pred_scores=pscores[pred_mask], max_detections=max_detections + ) + elif not np.any(pred_mask): # no predictions + result[c] = _matching_no_pred(iou_thresholds=iou_thresholds, gt_ignore=gignore[gt_mask]) + else: # at least one prediction and one ground truth + result[c] = _matching_single_image_single_class( + iou_fn=iou_fn, + pred_boxes=pboxes[pred_mask], + pred_scores=pscores[pred_mask], + gt_boxes=gboxes[gt_mask], + gt_ignore=gignore[gt_mask], + max_detections=max_detections, + iou_thresholds=iou_thresholds, + ) + results.append(result) + return results + + +def _matching_no_gt( + iou_thresholds: Sequence[float], pred_scores: np.ndarray, max_detections: int +) -> Dict[str, np.ndarray]: + """ + Matching result with not ground truth in image + + Args: + iou_thresholds: defined which IoU thresholds should be evaluated + dt_scores: predicted scores + max_detections: maximum number of allowed detections per image. + This functions uses this parameter to stay consistent with + the actual matching function which needs this limit. + + Returns: + computed matching, a Dict[str, np.ndarray] + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + """ + dt_ind = np.argsort(-pred_scores, kind="mergesort") + dt_ind = dt_ind[:max_detections] + dt_scores = pred_scores[dt_ind] + + num_preds = len(dt_scores) + + gt_match: np.ndarray = np.array([[]] * len(iou_thresholds)) + dt_match: np.ndarray = np.zeros((len(iou_thresholds), num_preds)) + dt_ignore: np.ndarray = np.zeros((len(iou_thresholds), num_preds)) + + return { + "dtMatches": dt_match, # [T, D], where T = number of thresholds, D = number of detections + "gtMatches": gt_match, # [T, G], where T = number of thresholds, G = number of ground truth + "dtScores": dt_scores, # [D] detection scores + "gtIgnore": np.array([]).reshape(-1), # [G] indicate whether ground truth should be ignored + "dtIgnore": dt_ignore, # [T, D], indicate which detections should be ignored + } + + +def _matching_no_pred(iou_thresholds: Sequence[float], gt_ignore: np.ndarray) -> Dict[str, np.ndarray]: + """ + Matching result with no predictions + + Args: + iou_thresholds: defined which IoU thresholds should be evaluated + gt_ignore: specified if which ground truth boxes are not counted as + true positives (detections which match theses boxes are not + counted as false positives either); [G], G number of ground truth + + Returns: + dict: computed matching + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + """ + dt_scores: np.ndarray = np.array([]) + dt_match: np.ndarray = np.array([[]] * len(iou_thresholds)) + dt_ignore: np.ndarray = np.array([[]] * len(iou_thresholds)) + + n_gt = 0 if gt_ignore.size == 0 else gt_ignore.shape[0] + gt_match = np.zeros((len(iou_thresholds), n_gt)) + + return { + "dtMatches": dt_match, # [T, D], where T = number of thresholds, D = number of detections + "gtMatches": gt_match, # [T, G], where T = number of thresholds, G = number of ground truth + "dtScores": dt_scores, # [D] detection scores + "gtIgnore": gt_ignore.reshape(-1), # [G] indicate whether ground truth should be ignored + "dtIgnore": dt_ignore, # [T, D], indicate which detections should be ignored + } + + +def _matching_single_image_single_class( + iou_fn: Callable[[np.ndarray, np.ndarray], np.ndarray], + pred_boxes: np.ndarray, + pred_scores: np.ndarray, + gt_boxes: np.ndarray, + gt_ignore: np.ndarray, + max_detections: int, + iou_thresholds: Sequence[float], +) -> Dict[str, np.ndarray]: + """ + Adapted from https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py + + Args: + iou_fn: compute overlap for each pair + iou_thresholds: defined which IoU thresholds should be evaluated + pred_boxes: predicted boxes from single batch; [D, dim * 2], D number + of predictions + pred_scores: predicted score for each bounding box; [D], D number of + predictions + gt_boxes: ground truth boxes; [G, dim * 2], G number of ground truth + gt_ignore: specified if which ground truth boxes are not counted as + true positives (detections which match theses boxes are not + counted as false positives either); [G], G number of ground truth + max_detections: maximum number of detections which should be evaluated + + Returns: + dict: computed matching + + - `dtMatches`: matched detections [T, D], where T = number of + thresholds, D = number of detections + - `gtMatches`: matched ground truth boxes [T, G], where T = number + of thresholds, G = number of ground truth + - `dtScores`: prediction scores [D] detection scores + - `gtIgnore`: ground truth boxes which should be ignored + [G] indicate whether ground truth should be ignored + - `dtIgnore`: detections which should be ignored [T, D], + indicate which detections should be ignored + """ + # filter for max_detections highest scoring predictions to speed up computation + dt_ind = np.argsort(-pred_scores, kind="mergesort") + dt_ind = dt_ind[:max_detections] + + pred_boxes = pred_boxes[dt_ind] + pred_scores = pred_scores[dt_ind] + + # sort ignored ground truth to last positions + gt_ind = np.argsort(gt_ignore, kind="mergesort") + gt_boxes = gt_boxes[gt_ind] + gt_ignore = gt_ignore[gt_ind] + + # ious between sorted(!) predictions and ground truth + ious = iou_fn(pred_boxes, gt_boxes) # array sized (num_preds, num_gts) + + num_preds, num_gts = ious.shape[0], ious.shape[1] + gt_match = np.zeros((len(iou_thresholds), num_gts)) + dt_match = np.zeros((len(iou_thresholds), num_preds)) + dt_ignore = np.zeros((len(iou_thresholds), num_preds)) + + for tind, t in enumerate(iou_thresholds): + for dind, _d in enumerate(pred_boxes): # iterate detections starting from highest scoring one + # information about best match so far (m=-1 -> unmatched) + iou = min([t, 1 - 1e-10]) + m = -1 + + for gind, _g in enumerate(gt_boxes): # iterate ground truth + # if this gt already matched, continue + if gt_match[tind, gind] > 0: + continue + + # if dt matched to reg gt, and on ignore gt, stop + if m > -1 and gt_ignore[m] == 0 and gt_ignore[gind] == 1: + break + + # continue to next gt unless better match made + if ious[dind, gind] < iou: + continue + + # if match successful and best so far, store appropriately + iou = ious[dind, gind] + m = gind + + # if match made, store id of match for both dt and gt + if m == -1: + continue + else: + dt_ignore[tind, dind] = int(gt_ignore[m]) + dt_match[tind, dind] = 1 + gt_match[tind, m] = 1 + + # store results for given image and category + return { + "dtMatches": dt_match, # [T, D], where T = number of thresholds, D = number of detections + "gtMatches": gt_match, # [T, G], where T = number of thresholds, G = number of ground truth + "dtScores": pred_scores, # [D] detection scores + "gtIgnore": gt_ignore.reshape(-1), # [G] indicate whether ground truth should be ignored + "dtIgnore": dt_ignore, # [T, D], indicate which detections should be ignored + } diff --git a/monai/apps/detection/networks/__init__.py b/monai/apps/detection/networks/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/networks/__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/monai/apps/detection/networks/retinanet_detector.py b/monai/apps/detection/networks/retinanet_detector.py new file mode 100644 index 0000000000..fd270ee094 --- /dev/null +++ b/monai/apps/detection/networks/retinanet_detector.py @@ -0,0 +1,1040 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE + +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +""" +Part of this script is adapted from +https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +""" + +import warnings +from typing import Any, Callable, Dict, List, Sequence, Tuple, Union + +import torch +from torch import Tensor, nn + +from monai.apps.detection.networks.retinanet_network import RetinaNet, resnet_fpn_feature_extractor +from monai.apps.detection.utils.anchor_utils import AnchorGenerator +from monai.apps.detection.utils.ATSS_matcher import ATSSMatcher +from monai.apps.detection.utils.box_coder import BoxCoder +from monai.apps.detection.utils.box_selector import BoxSelector +from monai.apps.detection.utils.detector_utils import check_training_targets, preprocess_images +from monai.apps.detection.utils.hard_negative_sampler import HardNegativeSampler +from monai.apps.detection.utils.predict_utils import ensure_dict_value_to_list_, predict_with_inferer +from monai.data.box_utils import box_iou +from monai.inferers import SlidingWindowInferer +from monai.networks.nets import resnet +from monai.utils import BlendMode, PytorchPadMode, ensure_tuple_rep, optional_import + +BalancedPositiveNegativeSampler, _ = optional_import( + "torchvision.models.detection._utils", name="BalancedPositiveNegativeSampler" +) +Matcher, _ = optional_import("torchvision.models.detection._utils", name="Matcher") + + +class RetinaNetDetector(nn.Module): + """ + Retinanet detector, expandable to other one stage anchor based box detectors in the future. + An example of construction can found in the source code of + :func:`~monai.apps.detection.networks.retinanet_detector.retinanet_resnet50_fpn_detector` . + + The input to the model is expected to be a list of tensors, each of shape (C, H, W) or (C, H, W, D), + one for each image, and should be in 0-1 range. Different images can have different sizes. + Or it can also be a Tensor sized (B, C, H, W) or (B, C, H, W, D). In this case, all images have same size. + + The behavior of the model changes depending if it is in training or evaluation mode. + + During training, the model expects both the input tensors, as well as a targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]`` or ``FloatTensor[N, 6]``): the ground-truth boxes in ``StandardMode``, i.e., + ``[xmin, ymin, xmax, ymax]`` or ``[xmin, ymin, zmin, xmax, ymax, zmax]`` format, + with ``0 <= xmin < xmax <= H``, ``0 <= ymin < ymax <= W``, ``0 <= zmin < zmax <= D``. + - labels: the class label for each ground-truth box + + The model returns a Dict[str, Tensor] during training, containing the classification and regression + losses. + When saving the model, only self.network contains trainable parameters and needs to be saved. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + + - boxes (``FloatTensor[N, 4]`` or ``FloatTensor[N, 6]``): the predicted boxes in ``StandardMode``, i.e., + ``[xmin, ymin, xmax, ymax]`` or ``[xmin, ymin, zmin, xmax, ymax, zmax]`` format, + with ``0 <= xmin < xmax <= H``, ``0 <= ymin < ymax <= W``, ``0 <= zmin < zmax <= D``. + - labels (Int64Tensor[N]): the predicted labels for each image + - labels_scores (Tensor[N]): the scores for each prediction + + Args: + network: a network that takes an image Tensor sized (B, C, H, W) or (B, C, H, W, D) as input + and outputs a dictionary Dict[str, List[Tensor]] or Dict[str, Tensor]. + anchor_generator: anchor generator. + box_overlap_metric: func that compute overlap between two sets of boxes, default is Intersection over Union (IoU). + debug: whether to print out internal parameters, used for debugging and parameter tuning. + + Notes: + + Input argument ``network`` can be a monai.apps.detection.networks.retinanet_network.RetinaNet(*) object, + but any network that meets the following rules is a valid input ``network``. + + 1. It should have attributes including spatial_dims, num_classes, cls_key, box_reg_key, num_anchors, size_divisible. + + - spatial_dims (int) is the spatial dimension of the network, we support both 2D and 3D. + - num_classes (int) is the number of classes, excluding the background. + - size_divisible (int or Sequene[int]) is the expection on the input image shape. + The network needs the input spatial_size to be divisible by size_divisible, length should be 2 or 3. + - cls_key (str) is the key to represent classification in the output dict. + - box_reg_key (str) is the key to represent box regression in the output dict. + - num_anchors (int) is the number of anchor shapes at each location. it should equal to + ``self.anchor_generator.num_anchors_per_location()[0]``. + + 2. Its input should be an image Tensor sized (B, C, H, W) or (B, C, H, W, D). + + 3. About its output ``head_outputs``: + + - It should be a dictionary with at least two keys: + ``network.cls_key`` and ``network.box_reg_key``. + - ``head_outputs[network.cls_key]`` should be List[Tensor] or Tensor. Each Tensor represents + classification logits map at one resolution level, + sized (B, num_classes*num_anchors, H_i, W_i) or (B, num_classes*num_anchors, H_i, W_i, D_i). + - ``head_outputs[network.box_reg_key]`` should be List[Tensor] or Tensor. Each Tensor represents + box regression map at one resolution level, + sized (B, 2*spatial_dims*num_anchors, H_i, W_i)or (B, 2*spatial_dims*num_anchors, H_i, W_i, D_i). + - ``len(head_outputs[network.cls_key]) == len(head_outputs[network.box_reg_key])``. + + Example: + + .. code-block:: python + + # define a naive network + import torch + class NaiveNet(torch.nn.Module): + def __init__(self, spatial_dims: int, num_classes: int): + super().__init__() + self.spatial_dims = spatial_dims + self.num_classes = num_classes + self.size_divisible = 2 + self.cls_key = "cls" + self.box_reg_key = "box_reg" + self.num_anchors = 1 + def forward(self, images: torch.Tensor): + spatial_size = images.shape[-self.spatial_dims:] + out_spatial_size = tuple(s//self.size_divisible for s in spatial_size) # half size of input + out_cls_shape = (images.shape[0],self.num_classes*self.num_anchors) + out_spatial_size + out_box_reg_shape = (images.shape[0],2*self.spatial_dims*self.num_anchors) + out_spatial_size + return {self.cls_key: [torch.randn(out_cls_shape)], self.box_reg_key: [torch.randn(out_box_reg_shape)]} + + # create a RetinaNetDetector detector + spatial_dims = 3 + num_classes = 5 + anchor_generator = monai.apps.detection.utils.anchor_utils.AnchorGeneratorWithAnchorShape( + feature_map_scales=(1, ), base_anchor_shapes=((8,) * spatial_dims) + ) + net = NaiveNet(spatial_dims, num_classes) + detector = RetinaNetDetector(net, anchor_generator) + + # only detector.network may contain trainable parameters. + optimizer = torch.optim.SGD( + detector.network.parameters(), + 1e-3, + momentum=0.9, + weight_decay=3e-5, + nesterov=True, + ) + torch.save(detector.network.state_dict(), 'model.pt') # save model + detector.network.load_state_dict(torch.load('model.pt')) # load model + """ + + def __init__( + self, network, anchor_generator: AnchorGenerator, box_overlap_metric: Callable = box_iou, debug: bool = False + ): + super().__init__() + + if not all( + hasattr(network, attr) + for attr in ["spatial_dims", "num_classes", "cls_key", "box_reg_key", "num_anchors", "size_divisible"] + ): + raise AttributeError( + "network should have attributes, including: " + "'spatial_dims', 'num_classes', 'cls_key', 'box_reg_key', 'num_anchors', 'size_divisible'." + ) + + self.network = network + self.spatial_dims = self.network.spatial_dims + self.num_classes = self.network.num_classes + self.size_divisible = ensure_tuple_rep(self.network.size_divisible, self.spatial_dims) + # keys for the network output + self.cls_key = self.network.cls_key + self.box_reg_key = self.network.box_reg_key + + # check if anchor_generator matches with network + self.anchor_generator = anchor_generator + + self.num_anchors_per_loc = self.anchor_generator.num_anchors_per_location()[0] + if self.num_anchors_per_loc != self.network.num_anchors: + raise ValueError( + f"Number of feature map channels ({self.network.num_anchors}) " + f"should match with number of anchors at each location ({self.num_anchors_per_loc})." + ) + # if new coming input images has same shape with + # self.previous_image_shape, there is no need to generate new anchors. + self.anchors: Union[List[Tensor], None] = None + self.previous_image_shape: Union[Any, None] = None + + self.box_overlap_metric = box_overlap_metric + self.debug = debug + + # default setting for training + self.fg_bg_sampler: Union[Any, None] = None + self.set_cls_loss(torch.nn.BCEWithLogitsLoss(reduction="mean")) # classification loss + self.set_box_regression_loss( + torch.nn.SmoothL1Loss(beta=1.0 / 9, reduction="mean"), encode_gt=True, decode_pred=False + ) # box regression loss + + # default setting for both training and inference + # can be updated by self.set_box_coder_weights(*) + self.box_coder = BoxCoder(weights=(1.0,) * 2 * self.spatial_dims) + + # default keys in the ground truth targets and predicted boxes, + # can be updated by self.set_target_keys(*) + self.target_box_key = "boxes" + self.target_label_key = "labels" + self.pred_score_key = self.target_label_key + "_scores" # score key for the detected boxes + + # default setting for inference, + # can be updated by self.set_sliding_window_inferer(*) + self.inferer: Union[SlidingWindowInferer, None] = None + # can be updated by self.set_box_selector_parameters(*), + self.box_selector = BoxSelector( + box_overlap_metric=self.box_overlap_metric, + score_thresh=0.05, + topk_candidates_per_level=1000, + nms_thresh=0.5, + detections_per_img=300, + apply_sigmoid=True, + ) + + def set_box_coder_weights(self, weights: Tuple[float]): + """ + Set the weights for box coder. + + Args: + weights: a list/tuple with length of 2*self.spatial_dims + + """ + if len(weights) != 2 * self.spatial_dims: + raise ValueError(f"len(weights) should be {2 * self.spatial_dims}, got weights={weights}.") + self.box_coder = BoxCoder(weights=weights) + + def set_target_keys(self, box_key: str, label_key: str): + """ + Set keys for the training targets and inference outputs. + During training, both box_key and label_key should be keys in the targets + when performing ``self.forward(input_images, targets)``. + During inference, they will be the keys in the output dict of `self.forward(input_images)``. + """ + self.target_box_key = box_key + self.target_label_key = label_key + self.pred_score_key = label_key + "_scores" + + def set_cls_loss(self, cls_loss: nn.Module) -> None: + """ + Using for training. Set loss for classification that takes logits as inputs, make sure sigmoid/softmax is built in. + + Args: + cls_loss: loss module for classification + + Example: + .. code-block:: python + + detector.set_cls_loss(torch.nn.BCEWithLogitsLoss(reduction="mean")) + detector.set_cls_loss(FocalLoss(reduction="mean", gamma=2.0)) + """ + self.cls_loss_func = cls_loss + + def set_box_regression_loss(self, box_loss: nn.Module, encode_gt: bool, decode_pred: bool) -> None: + """ + Using for training. Set loss for box regression. + + Args: + box_loss: loss module for box regression + encode_gt: if True, will encode ground truth boxes to target box regression + before computing the losses. Should be True for L1 loss and False for GIoU loss. + decode_pred: if True, will decode predicted box regression into predicted boxes + before computing losses. Should be False for L1 loss and True for GIoU loss. + + Example: + .. code-block:: python + + detector.set_box_regression_loss( + torch.nn.SmoothL1Loss(beta=1.0 / 9, reduction="mean"), + encode_gt = True, decode_pred = False + ) + detector.set_box_regression_loss( + monai.losses.giou_loss.BoxGIoULoss(reduction="mean"), + encode_gt = False, decode_pred = True + ) + """ + self.box_loss_func = box_loss + self.encode_gt = encode_gt + self.decode_pred = decode_pred + + def set_regular_matcher(self, fg_iou_thresh: float, bg_iou_thresh: float, allow_low_quality_matches=True) -> None: + """ + Using for training. Set torchvision matcher that matches anchors with ground truth boxes. + + Args: + fg_iou_thresh: foreground IoU threshold for Matcher, considered as matched if IoU > fg_iou_thresh + bg_iou_thresh: background IoU threshold for Matcher, considered as not matched if IoU < bg_iou_thresh + """ + if fg_iou_thresh < bg_iou_thresh: + raise ValueError( + "Require fg_iou_thresh >= bg_iou_thresh. " + f"Got fg_iou_thresh={fg_iou_thresh}, bg_iou_thresh={bg_iou_thresh}." + ) + self.proposal_matcher = Matcher(fg_iou_thresh, bg_iou_thresh, allow_low_quality_matches=True) + + def set_atss_matcher(self, num_candidates: int = 4, center_in_gt: bool = False) -> None: + """ + Using for training. Set ATSS matcher that matches anchors with ground truth boxes + + Args: + num_candidates: number of positions to select candidates from. + Smaller value will result in a higher matcher threshold and less matched candidates. + center_in_gt: If False (default), matched anchor center points do not need + to lie withing the ground truth box. Recommend False for small objects. + If True, will result in a strict matcher and less matched candidates. + """ + self.proposal_matcher = ATSSMatcher(num_candidates, self.box_overlap_metric, center_in_gt, debug=self.debug) + + def set_hard_negative_sampler( + self, batch_size_per_image: int, positive_fraction: float, min_neg: int = 1, pool_size: float = 10 + ): + """ + Using for training. Set hard negative sampler that samples part of the anchors for training. + + HardNegativeSampler is used to suppress false positive rate in classification tasks. + During training, it select negative samples with high prediction scores. + + Args: + batch_size_per_image: number of elements to be selected per image + positive_fraction: percentage of positive elements in the selected samples + min_neg: minimum number of negative samples to select if possible. + pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from + ``num_neg * pool_size`` negative samples with the highest prediction scores. + Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', + i.e., negative samples with lower prediction scores. + """ + self.fg_bg_sampler = HardNegativeSampler( + batch_size_per_image=batch_size_per_image, + positive_fraction=positive_fraction, + min_neg=min_neg, + pool_size=pool_size, + ) + + def set_balanced_sampler(self, batch_size_per_image: int, positive_fraction: float): + """ + Using for training. Set torchvision balanced sampler that samples part of the anchors for training. + + Args: + batch_size_per_image: number of elements to be selected per image + positive_fraction: percentage of positive elements per batch + + """ + self.fg_bg_sampler = BalancedPositiveNegativeSampler( + batch_size_per_image=batch_size_per_image, positive_fraction=positive_fraction + ) + + def set_sliding_window_inferer( + self, + roi_size: Union[Sequence[int], int], + sw_batch_size: int = 1, + overlap: float = 0.5, + mode: Union[BlendMode, str] = BlendMode.CONSTANT, + sigma_scale: Union[Sequence[float], float] = 0.125, + padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT, + cval: float = 0.0, + sw_device: Union[torch.device, str, None] = None, + device: Union[torch.device, str, None] = None, + progress: bool = False, + cache_roi_weight_map: bool = False, + ): + """ + Define sliding window inferer and store it to self.inferer. + """ + self.inferer = SlidingWindowInferer( + roi_size, + sw_batch_size, + overlap, + mode, + sigma_scale, + padding_mode, + cval, + sw_device, + device, + progress, + cache_roi_weight_map, + ) + + def set_box_selector_parameters( + self, + score_thresh: float = 0.05, + topk_candidates_per_level: int = 1000, + nms_thresh: float = 0.5, + detections_per_img: int = 300, + apply_sigmoid: bool = True, + ): + """ + Using for inference. Set the parameters that are used for box selection during inference. + The box selection is performed with the following steps: + + #. For each level, discard boxes with scores less than self.score_thresh. + #. For each level, keep boxes with top self.topk_candidates_per_level scores. + #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overapping threshold nms_thresh. + #. For the whole image, keep boxes with top self.detections_per_img scores. + + Args: + score_thresh: no box with scores less than score_thresh will be kept + topk_candidates_per_level: max number of boxes to keep for each level + nms_thresh: box overlapping threshold for NMS + detections_per_img: max number of boxes to keep for each image + """ + + self.box_selector = BoxSelector( + box_overlap_metric=self.box_overlap_metric, + apply_sigmoid=apply_sigmoid, + score_thresh=score_thresh, + topk_candidates_per_level=topk_candidates_per_level, + nms_thresh=nms_thresh, + detections_per_img=detections_per_img, + ) + + def forward( + self, + input_images: Union[List[Tensor], Tensor], + targets: Union[List[Dict[str, Tensor]], None] = None, + use_inferer: bool = False, + ) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]: + """ + Returns a dict of losses during training, or a list predicted dict of boxes and labels during inference. + + Args: + input_images: The input to the model is expected to be a list of tensors, each of shape (C, H, W) or (C, H, W, D), + one for each image, and should be in 0-1 range. Different images can have different sizes. + Or it can also be a Tensor sized (B, C, H, W) or (B, C, H, W, D). In this case, all images have same size. + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image (optional). + use_inferer: whether to use self.inferer, a sliding window inferer, to do the inference. + If False, will simply forward the network. + If True, will use self.inferer, and requires + ``self.set_sliding_window_inferer(*args)`` to have been called before. + + Return: + If training mode, will return a dict with at least two keys, + including self.cls_key and self.box_reg_key, representing classification loss and box regression loss. + + If evaluation mode, will return a list of detection results. + Each element corresponds to an images in ``input_images``, is a dict with at least three keys, + including self.target_box_key, self.target_label_key, self.pred_score_key, + representing predicted boxes, classification labels, and classification scores. + + """ + # 1. Check if input arguments are valid + if self.training: + check_training_targets(input_images, targets, self.spatial_dims, self.target_label_key, self.target_box_key) + self._check_detector_training_components() + + # 2. Pad list of images to a single Tensor `images` with spatial size divisible by self.size_divisible. + # image_sizes stores the original spatial_size of each image before padding. + images, image_sizes = preprocess_images(input_images, self.spatial_dims, self.size_divisible) + + # 3. Generate network outputs. Use inferer only in evaluation mode. + if self.training or (not use_inferer): + head_outputs = self.network(images) + ensure_dict_value_to_list_(head_outputs) # ensure head_outputs is Dict[str, List[Tensor]] + else: + if self.inferer is None: + raise ValueError( + "`self.inferer` is not defined." "Please refer to function self.set_sliding_window_inferer(*)." + ) + head_outputs = predict_with_inferer( + images, self.network, keys=[self.cls_key, self.box_reg_key], inferer=self.inferer + ) + + # 4. Generate anchors and store it in self.anchors: List[Tensor] + 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]] + + # 5. Reshape and concatenate head_outputs values from List[Tensor] to Tensor + # head_outputs, originally being Dict[str, List[Tensor]], will be reshaped to Dict[str, Tensor] + for key in [self.cls_key, self.box_reg_key]: + # 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 + head_outputs[key] = self._reshape_maps(head_outputs[key]) + + # 6(1). If during training, return losses + if self.training: + losses = self.compute_loss(head_outputs, targets, self.anchors, num_anchor_locs_per_level) # type: ignore + return losses + + # 6(2). If during inference, return detection results + detections = self.postprocess_detections( + head_outputs, self.anchors, image_sizes, num_anchor_locs_per_level # type: ignore + ) + return detections + + def _check_detector_training_components(self): + """ + Check if self.proposal_matcher and self.fg_bg_sampler have been set for training. + """ + if not hasattr(self, "proposal_matcher"): + raise AttributeError( + "Matcher is not set. Please refer to self.set_regular_matcher(*) or self.set_atss_matcher(*)." + ) + if self.fg_bg_sampler is None and self.debug: + warnings.warn( + "No balanced sampler is used. Negative samples are likely to " + "be much more than positive samples. Please set balanced samplers with self.set_balanced_sampler(*) " + "or self.set_hard_negative_sampler(*), " + "or set classification loss function as Focal loss with self.set_cls_loss(*)" + ) + + def generate_anchors(self, images: Tensor, head_outputs: Dict[str, List[Tensor]]): + """ + Generate anchors and store it in self.anchors: List[Tensor]. + We generate anchors only when there is no stored anchors, + or the new coming images has different shape with self.previous_image_shape + + Args: + images: input images, a (B, C, H, W) or (B, C, H, W, D) Tensor. + head_outputs: head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor + sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor + sized (B, sum(HW(D)A), 2*self.spatial_dims) + """ + if (self.anchors is None) or (self.previous_image_shape != images.shape): + self.anchors = self.anchor_generator(images, head_outputs[self.cls_key]) # List[Tensor], len = batchsize + self.previous_image_shape = images.shape + + def _reshape_maps(self, result_maps: List[Tensor]) -> Tensor: + """ + Concat network output map list to a single Tensor. + This function is used in both training and inference. + + Args: + result_maps: a list of Tensor, each Tensor is a (B, num_channel*A, H, W) or (B, num_channel*A, H, W, D) map. + A = self.num_anchors_per_loc + + Return: + reshaped and concatenated result, sized (B, sum(HWA), num_channel) or (B, sum(HWDA), num_channel) + """ + all_reshaped_result_map = [] + + for result_map in result_maps: + batch_size = result_map.shape[0] + num_channel = result_map.shape[1] // self.num_anchors_per_loc + spatial_size = result_map.shape[-self.spatial_dims :] + + # reshaped_result_map will become (B, A, num_channel, H, W) or (B, A, num_channel, H, W, D) + # A = self.num_anchors_per_loc + view_shape = (batch_size, -1, num_channel) + spatial_size + reshaped_result_map = result_map.view(view_shape) + + # permute output to (B, H, W, A, num_channel) or (B, H, W, D, A, num_channel) + if self.spatial_dims == 2: + reshaped_result_map = reshaped_result_map.permute(0, 3, 4, 1, 2) + elif self.spatial_dims == 3: + reshaped_result_map = reshaped_result_map.permute(0, 3, 4, 5, 1, 2) + else: + ValueError("Images can only be 2D or 3D.") + + # reshaped_result_map will become (B, HWA, num_channel) or (B, HWDA, num_channel) + reshaped_result_map = reshaped_result_map.reshape(batch_size, -1, num_channel) + + if torch.isnan(reshaped_result_map).any() or torch.isinf(reshaped_result_map).any(): + raise ValueError("Concatenated result is NaN or Inf.") + + all_reshaped_result_map.append(reshaped_result_map) + + return torch.cat(all_reshaped_result_map, dim=1) + + def postprocess_detections( + self, + head_outputs_reshape: Dict[str, Tensor], + anchors: List[Tensor], + image_sizes: List[List[int]], + num_anchor_locs_per_level: Sequence[int], + need_sigmoid: bool = True, + ) -> List[Dict[str, Tensor]]: + """ + Postprocessing to generate detection result from classification logits and box regression. + Use self.box_selector to select the final outut boxes for each image. + + Args: + head_outputs_reshape: reshaped head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor + sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor + sized (B, sum(HW(D)A), 2*self.spatial_dims) + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + anchors: a list of Tensor. Each Tensor represents anchors for each image, + sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). + A = self.num_anchors_per_loc. + + Return: + a list of dict, each dict scorresponds to detection result on image. + """ + + # recover level sizes, HWA or HWDA for each level + num_anchors_per_level = [ + num_anchor_locs * self.num_anchors_per_loc for num_anchor_locs in num_anchor_locs_per_level + ] + + # split outputs per level + split_head_outputs: Dict[str, List[Tensor]] = {} + for k in head_outputs_reshape: + split_head_outputs[k] = list(head_outputs_reshape[k].split(num_anchors_per_level, dim=1)) + split_anchors = [list(a.split(num_anchors_per_level)) for a in anchors] # List[List[Tensor]] + + class_logits = split_head_outputs[self.cls_key] # List[Tensor], each sized (B, HWA, self.num_classes) + box_regression = split_head_outputs[self.box_reg_key] # List[Tensor], each sized (B, HWA, 2*spatial_dims) + compute_dtype = class_logits[0].dtype + + num_images = len(image_sizes) # B + + detections: List[Dict[str, Tensor]] = [] + + for index in range(num_images): + box_regression_per_image = [ + br[index] for br in box_regression + ] # List[Tensor], each sized (HWA, 2*spatial_dims) + logits_per_image = [cl[index] for cl in class_logits] # List[Tensor], each sized (HWA, self.num_classes) + anchors_per_image, img_spatial_size = split_anchors[index], image_sizes[index] + # decode box regression into boxes + boxes_per_image = [ + self.box_coder.decode_single(b.to(torch.float32), a).to(compute_dtype) + for b, a in zip(box_regression_per_image, anchors_per_image) + ] # List[Tensor], each sized (HWA, 2*spatial_dims) + + selected_boxes, selected_scores, selected_labels = self.box_selector.select_boxes_per_image( + boxes_per_image, logits_per_image, img_spatial_size + ) + + detections.append( + { + self.target_box_key: selected_boxes, # Tensor, sized (N, 2*spatial_dims) + self.pred_score_key: selected_scores, # Tensor, sized (N, ) + self.target_label_key: selected_labels, # Tensor, sized (N, ) + } + ) + + return detections + + def compute_loss( + self, + head_outputs_reshape: Dict[str, Tensor], + targets: List[Dict[str, Tensor]], + anchors: List[Tensor], + num_anchor_locs_per_level: Sequence[int], + ) -> Dict[str, Tensor]: + """ + Compute losses. + + Args: + head_outputs_reshape: reshaped head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor + sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor + sized (B, sum(HW(D)A), 2*self.spatial_dims) + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + anchors: a list of Tensor. Each Tensor represents anchors for each image, + sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). + A = self.num_anchors_per_loc. + + Return: + a dict of several kinds of losses. + """ + matched_idxs = self.compute_anchor_matched_idxs(anchors, targets, num_anchor_locs_per_level) + losses_cls = self.compute_cls_loss(head_outputs_reshape[self.cls_key], targets, matched_idxs) + losses_box_regression = self.compute_box_loss( + head_outputs_reshape[self.box_reg_key], targets, anchors, matched_idxs + ) + return {self.cls_key: losses_cls, self.box_reg_key: losses_box_regression} + + def compute_anchor_matched_idxs( + self, anchors: List[Tensor], targets: List[Dict[str, Tensor]], num_anchor_locs_per_level: Sequence[int] + ) -> List[Tensor]: + """ + Compute the matched indices between anchors and ground truth (gt) boxes in targets. + output[k][i] represents the matched gt index for anchor[i] in image k. + Suppose there are M gt boxes for image k. The range of it output[k][i] value is [-2, -1, 0, ..., M-1]. + [0, M - 1] indicates this anchor is matched with a gt box, + while a negative value indicating that it is not matched. + + Args: + anchors: a list of Tensor. Each Tensor represents anchors for each image, + sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). + A = self.num_anchors_per_loc. + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + num_anchor_locs_per_level: each element represents HW or HWD at this level. + + + Return: + a list of matched index `matched_idxs_per_image` (Tensor[int64]), Tensor sized (sum(HWA),) or (sum(HWDA),). + Suppose there are M gt boxes. `matched_idxs_per_image[i]` is a matched gt index in [0, M - 1] + or a negative value indicating that anchor i could not be matched. + BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 + """ + matched_idxs = [] + for anchors_per_image, targets_per_image in zip(anchors, targets): + # anchors_per_image: Tensor, targets_per_image: Dice[str, Tensor] + if targets_per_image[self.target_box_key].numel() == 0: + # if no GT boxes + matched_idxs.append( + torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64, device=anchors_per_image.device) + ) + continue + + # matched_idxs_per_image (Tensor[int64]): Tensor sized (sum(HWA),) or (sum(HWDA),) + # Suppose there are M gt boxes. matched_idxs_per_image[i] is a matched gt index in [0, M - 1] + # or a negative value indicating that anchor i could not be matched. + # BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 + if isinstance(self.proposal_matcher, Matcher): + # if torcvision matcher + match_quality_matrix = self.box_overlap_metric( + targets_per_image[self.target_box_key].to(anchors_per_image.device), anchors_per_image + ) + matched_idxs_per_image = self.proposal_matcher(match_quality_matrix) + elif isinstance(self.proposal_matcher, ATSSMatcher): + # if monai ATSS matcher + match_quality_matrix, matched_idxs_per_image = self.proposal_matcher( + targets_per_image[self.target_box_key].to(anchors_per_image.device), + anchors_per_image, + num_anchor_locs_per_level, + self.num_anchors_per_loc, + ) + else: + raise NotImplementedError( + "Currently support torchvision Matcher and monai ATSS matcher. Other types of matcher not supported. " + "Please override self.compute_anchor_matched_idxs(*) for your own matcher." + ) + + if self.debug: + print(f"Max box overlap between anchors and gt boxes: {torch.max(match_quality_matrix,dim=1)[0]}.") + + if torch.max(matched_idxs_per_image) < 0: + warnings.warn( + f"No anchor is matched with GT boxes. Please adjust matcher setting, anchor setting," + " or the network setting to change zoom scale between network output and input images." + f"GT boxes are {targets_per_image[self.target_box_key]}." + ) + + matched_idxs.append(matched_idxs_per_image) + return matched_idxs + + def compute_cls_loss( + self, cls_logits: Tensor, targets: List[Dict[str, Tensor]], matched_idxs: List[Tensor] + ) -> Tensor: + """ + Compute classification losses. + + Args: + cls_logits: classification logits, sized (B, sum(HW(D)A), self.num_classes) + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + matched_idxs: a list of matched index. each element is sized (sum(HWA),) or (sum(HWDA),) + + Return: + classification losses. + """ + total_cls_logits_list = [] + total_gt_classes_target_list = [] + for targets_per_image, cls_logits_per_image, matched_idxs_per_image in zip(targets, cls_logits, matched_idxs): + # for each image, get training samples + sampled_cls_logits_per_image, sampled_gt_classes_target = self.get_cls_train_sample_per_image( + cls_logits_per_image, targets_per_image, matched_idxs_per_image + ) + total_cls_logits_list.append(sampled_cls_logits_per_image) + total_gt_classes_target_list.append(sampled_gt_classes_target) + + total_cls_logits = torch.cat(total_cls_logits_list, dim=0) + total_gt_classes_target = torch.cat(total_gt_classes_target_list, dim=0) + losses: Tensor = self.cls_loss_func(total_cls_logits, total_gt_classes_target).to(total_cls_logits.dtype) + return losses + + def compute_box_loss( + self, + box_regression: Tensor, + targets: List[Dict[str, Tensor]], + anchors: List[Tensor], + matched_idxs: List[Tensor], + ) -> Tensor: + """ + Compute box regression losses. + + Args: + box_regression: box regression results, sized (B, sum(HWA), 2*self.spatial_dims) + targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + anchors: a list of Tensor. Each Tensor represents anchors for each image, + sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). + A = self.num_anchors_per_loc. + matched_idxs: a list of matched index. each element is sized (sum(HWA),) or (sum(HWDA),) + + Return: + box regression losses. + """ + total_box_regression_list = [] + total_target_regression_list = [] + + for targets_per_image, box_regression_per_image, anchors_per_image, matched_idxs_per_image in zip( + targets, box_regression, anchors, matched_idxs + ): + # for each image, get training samples + decode_box_regression_per_image, matched_gt_boxes_per_image = self.get_box_train_sample_per_image( + box_regression_per_image, targets_per_image, anchors_per_image, matched_idxs_per_image + ) + total_box_regression_list.append(decode_box_regression_per_image) + total_target_regression_list.append(matched_gt_boxes_per_image) + + total_box_regression = torch.cat(total_box_regression_list, dim=0) + total_target_regression = torch.cat(total_target_regression_list, dim=0) + + if total_box_regression.shape[0] == 0: + # if there is no training sample. + losses = torch.tensor(0.0) + return losses + + losses = self.box_loss_func(total_box_regression, total_target_regression).to(total_box_regression.dtype) + + return losses + + def get_cls_train_sample_per_image( + self, cls_logits_per_image: Tensor, targets_per_image: Dict[str, Tensor], matched_idxs_per_image: Tensor + ) -> Tuple[Tensor, Tensor]: + """ + Get samples from one image for classification losses computation. + + Args: + cls_logits_per_image: classification logits for one image, (sum(HWA), self.num_classes) + targets_per_image: a dict with at least two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + matched_idxs_per_image: matched index, Tensor sized (sum(HWA),) or (sum(HWDA),) + Suppose there are M gt boxes. matched_idxs_per_image[i] is a matched gt index in [0, M - 1] + or a negative value indicating that anchor i could not be matched. + BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 + + Return: + paired predicted and GT samples from one image for classification losses computation + """ + + if torch.isnan(cls_logits_per_image).any() or torch.isinf(cls_logits_per_image).any(): + raise ValueError("NaN or Inf in predicted classification logits.") + + foreground_idxs_per_image = matched_idxs_per_image >= 0 + + num_foreground = foreground_idxs_per_image.sum() + num_gt_box = targets_per_image[self.target_box_key].shape[0] + + if self.debug: + print(f"Number of positive (matched) anchors: {num_foreground}; Number of GT box: {num_gt_box}.") + if num_gt_box > 0 and num_foreground < 2 * num_gt_box: + print( + f"Only {num_foreground} anchors are matched with {num_gt_box} GT boxes. " + "Please consider adjusting matcher setting, anchor setting," + " or the network setting to change zoom scale between network output and input images." + ) + + # create the target classification with one-hot encoding + gt_classes_target = torch.zeros_like(cls_logits_per_image) # (sum(HW(D)A), self.num_classes) + gt_classes_target[ + foreground_idxs_per_image, # fg anchor idx in + targets_per_image[self.target_label_key][ + matched_idxs_per_image[foreground_idxs_per_image] + ], # fg class label + ] = 1.0 + + if self.fg_bg_sampler is None: + # if no balanced sampling + valid_idxs_per_image = matched_idxs_per_image != self.proposal_matcher.BETWEEN_THRESHOLDS + else: + # The input of fg_bg_sampler: list of tensors containing -1, 0 or positive values. + # Each tensor corresponds to a specific image. + # -1 values are ignored, 0 are considered as negatives and > 0 as positives. + + # matched_idxs_per_image (Tensor[int64]): an N tensor where N[i] is a matched gt in + # [0, M - 1] or a negative value indicating that prediction i could not + # be matched. BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 + if isinstance(self.fg_bg_sampler, HardNegativeSampler): + max_cls_logits_per_image = torch.max(cls_logits_per_image.to(torch.float32), dim=1)[0] + sampled_pos_inds_list, sampled_neg_inds_list = self.fg_bg_sampler( + [matched_idxs_per_image + 1], max_cls_logits_per_image + ) + elif isinstance(self.fg_bg_sampler, BalancedPositiveNegativeSampler): + sampled_pos_inds_list, sampled_neg_inds_list = self.fg_bg_sampler([matched_idxs_per_image + 1]) + else: + raise NotImplementedError( + "Currently support torchvision BalancedPositiveNegativeSampler and monai HardNegativeSampler matcher. " + "Other types of sampler not supported. " + "Please override self.get_cls_train_sample_per_image(*) for your own sampler." + ) + + sampled_pos_inds = torch.where(torch.cat(sampled_pos_inds_list, dim=0))[0] + sampled_neg_inds = torch.where(torch.cat(sampled_neg_inds_list, dim=0))[0] + valid_idxs_per_image = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0) + + return cls_logits_per_image[valid_idxs_per_image, :], gt_classes_target[valid_idxs_per_image, :] + + def get_box_train_sample_per_image( + self, + box_regression_per_image: Tensor, + targets_per_image: Dict[str, Tensor], + anchors_per_image: Tensor, + matched_idxs_per_image: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Get samples from one image for box regression losses computation. + + Args: + box_regression_per_image: box regression result for one image, (sum(HWA), 2*self.spatial_dims) + targets_per_image: a dict with at least two keys: self.target_box_key and self.target_label_key, + ground-truth boxes present in the image. + anchors_per_image: anchors of one image, + sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). + A = self.num_anchors_per_loc. + matched_idxs_per_image: matched index, sized (sum(HWA),) or (sum(HWDA),) + + Return: + paired predicted and GT samples from one image for box regression losses computation + """ + + if torch.isnan(box_regression_per_image).any() or torch.isinf(box_regression_per_image).any(): + raise ValueError("NaN or Inf in predicted box regression.") + + foreground_idxs_per_image = torch.where(matched_idxs_per_image >= 0)[0] + num_gt_box = targets_per_image[self.target_box_key].shape[0] + + # if no GT box, return empty arrays + if num_gt_box == 0: + return box_regression_per_image[0:0, :], box_regression_per_image[0:0, :] + + # select only the foreground boxes + # matched GT boxes for foreground anchors + matched_gt_boxes_per_image = targets_per_image[self.target_box_key][ + matched_idxs_per_image[foreground_idxs_per_image] + ].to(box_regression_per_image.device) + # predicted box regression for foreground anchors + box_regression_per_image = box_regression_per_image[foreground_idxs_per_image, :] + # foreground anchors + anchors_per_image = anchors_per_image[foreground_idxs_per_image, :] + + # encode GT boxes or decode predicted box regression before computing losses + matched_gt_boxes_per_image_ = matched_gt_boxes_per_image + box_regression_per_image_ = box_regression_per_image + if self.encode_gt: + matched_gt_boxes_per_image_ = self.box_coder.encode_single(matched_gt_boxes_per_image_, anchors_per_image) + if self.decode_pred: + box_regression_per_image_ = self.box_coder.decode_single(box_regression_per_image_, anchors_per_image) + + return box_regression_per_image_, matched_gt_boxes_per_image_ + + +def retinanet_resnet50_fpn_detector( + num_classes: int, + anchor_generator: AnchorGenerator, + returned_layers: Sequence[int] = (1, 2, 3), + pretrained: bool = False, + progress: bool = True, + **kwargs: Any, +) -> RetinaNetDetector: + """ + Returns a RetinaNet detector using a ResNet-50 as backbone, which can be pretrained + from `Med3D: Transfer Learning for 3D Medical Image Analysis ` + _. + + Args: + num_classes: number of output classes of the model (excluding the background). + anchor_generator: AnchorGenerator, + returned_layers: returned layers to extract feature maps. Each returned layer should be in the range [1,4]. + len(returned_layers)+1 will be the number of extracted feature maps. + There is an extra maxpooling layer LastLevelMaxPool() appended. + pretrained: If True, returns a backbone pre-trained on 23 medical datasets + progress: If True, displays a progress bar of the download to stderr + + Return: + A RetinaNetDetector object with resnet50 as backbone + + Example: + + .. code-block:: python + + # define a naive network + resnet_param = { + "pretrained": False, + "spatial_dims": 3, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": 7, + "conv1_t_stride": (2, 2, 2) + } + returned_layers = [1] + anchor_generator = monai.apps.detection.utils.anchor_utils.AnchorGeneratorWithAnchorShape( + feature_map_scales=(1, 2), base_anchor_shapes=((8,) * resnet_param["spatial_dims"]) + ) + detector = retinanet_resnet50_fpn_detector( + **resnet_param, anchor_generator=anchor_generator, returned_layers=returned_layers + ) + """ + + backbone = resnet.resnet50(pretrained, progress, **kwargs) + spatial_dims = len(backbone.conv1.stride) + # number of output feature maps is len(returned_layers)+1 + feature_extractor = resnet_fpn_feature_extractor( + backbone=backbone, + spatial_dims=spatial_dims, + pretrained_backbone=pretrained, + trainable_backbone_layers=None, + returned_layers=returned_layers, + ) + num_anchors = anchor_generator.num_anchors_per_location()[0] + size_divisible = [s * 2 * 2 ** max(returned_layers) for s in feature_extractor.body.conv1.stride] + network = RetinaNet( + spatial_dims=spatial_dims, + num_classes=num_classes, + num_anchors=num_anchors, + feature_extractor=feature_extractor, + size_divisible=size_divisible, + ) + return RetinaNetDetector(network, anchor_generator) diff --git a/monai/apps/detection/networks/retinanet_network.py b/monai/apps/detection/networks/retinanet_network.py new file mode 100644 index 0000000000..4539a913ac --- /dev/null +++ b/monai/apps/detection/networks/retinanet_network.py @@ -0,0 +1,407 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE + +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +""" +Part of this script is adapted from +https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +""" + +import math +from typing import Callable, Dict, List, Sequence, Union + +import torch +from torch import Tensor, nn + +from monai.networks.blocks.backbone_fpn_utils import _resnet_fpn_extractor +from monai.networks.layers.factories import Conv +from monai.networks.nets import resnet +from monai.utils import ensure_tuple_rep, look_up_option, optional_import + +_validate_trainable_layers, _ = optional_import( + "torchvision.models.detection.backbone_utils", name="_validate_trainable_layers" +) + + +class RetinaNetClassificationHead(nn.Module): + """ + A classification head for use in RetinaNet. + + This head takes a list of feature maps as inputs, and outputs a list of classification maps. + Each output map has same spatial size with the corresponding input feature map, + and the number of output channel is num_anchors * num_classes. + + Args: + in_channels: number of channels of the input feature + num_anchors: number of anchors to be predicted + num_classes: number of classes to be predicted + spatial_dims: spatial dimension of the network, should be 2 or 3. + prior_probability: prior probability to initialize classification convolutional layers. + """ + + def __init__( + self, in_channels: int, num_anchors: int, num_classes: int, spatial_dims: int, prior_probability: float = 0.01 + ): + super().__init__() + + conv_type: Callable = Conv[Conv.CONV, spatial_dims] + conv = [] + for _ in range(4): + conv.append(conv_type(in_channels, in_channels, kernel_size=3, stride=1, padding=1)) + conv.append(nn.GroupNorm(num_groups=8, num_channels=in_channels)) + conv.append(nn.ReLU()) + self.conv = nn.Sequential(*conv) + + for layer in self.conv.children(): + if isinstance(layer, conv_type): # type: ignore + torch.nn.init.normal_(layer.weight, std=0.01) # type: ignore + torch.nn.init.constant_(layer.bias, 0) # type: ignore + + self.cls_logits = conv_type(in_channels, num_anchors * num_classes, kernel_size=3, stride=1, padding=1) + torch.nn.init.normal_(self.cls_logits.weight, std=0.01) + torch.nn.init.constant_(self.cls_logits.bias, -math.log((1 - prior_probability) / prior_probability)) + + self.num_classes = num_classes + self.num_anchors = num_anchors + + def forward(self, x: List[Tensor]) -> List[Tensor]: + """ + It takes a list of feature maps as inputs, and outputs a list of classification maps. + Each output classification map has same spatial size with the corresponding input feature map, + and the number of output channel is num_anchors * num_classes. + + Args: + x: list of feature map, x[i] is a (B, in_channels, H_i, W_i) or (B, in_channels, H_i, W_i, D_i) Tensor. + + Return: + cls_logits_maps, list of classification map. cls_logits_maps[i] is a + (B, num_anchors * num_classes, H_i, W_i) or (B, num_anchors * num_classes, H_i, W_i, D_i) Tensor. + + """ + cls_logits_maps = [] + + if isinstance(x, Tensor): + feature_maps = [x] + else: + feature_maps = x + + for features in feature_maps: + cls_logits = self.conv(features) + cls_logits = self.cls_logits(cls_logits) + + cls_logits_maps.append(cls_logits) + + if torch.isnan(cls_logits).any() or torch.isinf(cls_logits).any(): + raise ValueError("cls_logits is NaN or Inf.") + + return cls_logits_maps + + +class RetinaNetRegressionHead(nn.Module): + """ + A regression head for use in RetinaNet. + + This head takes a list of feature maps as inputs, and outputs a list of box regression maps. + Each output box regression map has same spatial size with the corresponding input feature map, + and the number of output channel is num_anchors * 2 * spatial_dims. + + Args: + in_channels: number of channels of the input feature + num_anchors: number of anchors to be predicted + spatial_dims: spatial dimension of the network, should be 2 or 3. + """ + + def __init__(self, in_channels: int, num_anchors: int, spatial_dims: int): + super().__init__() + + conv_type: Callable = Conv[Conv.CONV, spatial_dims] + + conv = [] + for _ in range(4): + conv.append(conv_type(in_channels, in_channels, kernel_size=3, stride=1, padding=1)) + conv.append(nn.GroupNorm(num_groups=8, num_channels=in_channels)) + conv.append(nn.ReLU()) + + self.conv = nn.Sequential(*conv) + + self.bbox_reg = conv_type(in_channels, num_anchors * 2 * spatial_dims, kernel_size=3, stride=1, padding=1) + torch.nn.init.normal_(self.bbox_reg.weight, std=0.01) + torch.nn.init.zeros_(self.bbox_reg.bias) + + for layer in self.conv.children(): + if isinstance(layer, conv_type): # type: ignore + torch.nn.init.normal_(layer.weight, std=0.01) # type: ignore + torch.nn.init.zeros_(layer.bias) # type: ignore + + def forward(self, x: List[Tensor]) -> List[Tensor]: + """ + It takes a list of feature maps as inputs, and outputs a list of box regression maps. + Each output box regression map has same spatial size with the corresponding input feature map, + and the number of output channel is num_anchors * 2 * spatial_dims. + + Args: + x: list of feature map, x[i] is a (B, in_channels, H_i, W_i) or (B, in_channels, H_i, W_i, D_i) Tensor. + + Return: + box_regression_maps, list of box regression map. cls_logits_maps[i] is a + (B, num_anchors * 2 * spatial_dims, H_i, W_i) or (B, num_anchors * 2 * spatial_dims, H_i, W_i, D_i) Tensor. + + """ + box_regression_maps = [] + + if isinstance(x, Tensor): + feature_maps = [x] + else: + feature_maps = x + + for features in feature_maps: + box_regression = self.conv(features) + box_regression = self.bbox_reg(box_regression) + + box_regression_maps.append(box_regression) + + if torch.isnan(box_regression).any() or torch.isinf(box_regression).any(): + raise ValueError("box_regression is NaN or Inf.") + + return box_regression_maps + + +class RetinaNet(nn.Module): + """ + The network used in RetinaNet. + + It takes an image tensor as inputs, and outputs a dictionary ``head_outputs``. + ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. + ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. + + Args: + spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images. + num_classes: number of output classes of the model (excluding the background). + num_anchors: number of anchors at each location. + feature_extractor: a network that outputs feature maps from the input images, + each feature map corresponds to a different resolution. + Its output can have a format of Tensor, Dict[Any, Tensor], or Sequence[Tensor]. + It can be the output of ``resnet_fpn_feature_extractor(*args, **kwargs)``. + size_divisible: the spatial size of the network input should be divisible by size_divisible, + decided by the feature_extractor. + + Example: + + .. code-block:: python + + from monai.networks.nets import resnet + spatial_dims = 3 # 3D network + conv1_t_stride = (2,2,1) # stride of first convolutional layer in backbone + backbone = resnet.ResNet( + spatial_dims = spatial_dims, + block = resnet.ResNetBottleneck, + layers = [3, 4, 6, 3], + block_inplanes = resnet.get_inplanes(), + n_input_channels= 1, + conv1_t_stride = conv1_t_stride, + conv1_t_size = (7,7,7), + ) + # This feature_extractor outputs 4-level feature maps. + # number of output feature maps is len(returned_layers)+1 + returned_layers = [1,2,3] # returned layer from feature pyramid network + feature_extractor = resnet_fpn_feature_extractor( + backbone = backbone, + spatial_dims = spatial_dims, + pretrained_backbone = False, + trainable_backbone_layers = None, + returned_layers = returned_layers, + ) + # This feature_extractor requires input image spatial size + # to be divisible by (32, 32, 16). + size_divisible = tuple(2*s*2**max(returned_layers) for s in conv1_t_stride) + model = RetinaNet( + spatial_dims = spatial_dims, + num_classes = 5, + num_anchors = 6, + feature_extractor=feature_extractor, + size_divisible = size_divisible, + ).to(device) + result = model(torch.rand(2, 1, 128,128,128)) + cls_logits_maps = result["cls_logits"] # a list of len(returned_layers)+1 Tensor + box_regression_maps = result["box_regression"] # a list of len(returned_layers)+1 Tensor + """ + + def __init__( + self, + spatial_dims: int, + num_classes: int, + num_anchors: int, + feature_extractor, + size_divisible: Union[Sequence[int], int] = 1, + ): + super().__init__() + + self.spatial_dims = look_up_option(spatial_dims, supported=[1, 2, 3]) + self.num_classes = num_classes + self.size_divisible = ensure_tuple_rep(size_divisible, self.spatial_dims) + + if not hasattr(feature_extractor, "out_channels"): + raise ValueError( + "feature_extractor should contain an attribute out_channels " + "specifying the number of output channels (assumed to be the " + "same for all the levels)" + ) + self.feature_extractor = feature_extractor + + self.feature_map_channels: int = self.feature_extractor.out_channels + self.num_anchors = num_anchors + self.classification_head = RetinaNetClassificationHead( + self.feature_map_channels, self.num_anchors, self.num_classes, spatial_dims=self.spatial_dims + ) + self.regression_head = RetinaNetRegressionHead( + self.feature_map_channels, self.num_anchors, spatial_dims=self.spatial_dims + ) + + self.cls_key: str = "classification" + self.box_reg_key: str = "box_regression" + + def forward(self, images: Tensor) -> Dict[str, List[Tensor]]: + """ + It takes an image tensor as inputs, and outputs a dictionary ``head_outputs``. + ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. + ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. + + Args: + images: input images, sized (B, img_channels, H, W) or (B, img_channels, H, W, D). + + Return: + a dictionary ``head_outputs`` with keys including self.cls_key and self.box_reg_key. + ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. + ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. + + """ + # compute features maps list from the input images. + features = self.feature_extractor(images) + if isinstance(features, Tensor): + feature_maps = [features] + elif torch.jit.isinstance(features, Dict[str, Tensor]): + feature_maps = list(features.values()) + else: + feature_maps = list(features) + + if not isinstance(feature_maps[0], Tensor): + raise ValueError("feature_extractor output format must be Tensor, Dict[str, Tensor], or Sequence[Tensor].") + + # compute classification and box regression maps from the feature maps + # expandable for mask prediction in the future + + head_outputs: Dict[str, List[Tensor]] = {self.cls_key: self.classification_head(feature_maps)} + head_outputs[self.box_reg_key] = self.regression_head(feature_maps) + + return head_outputs + + +def resnet_fpn_feature_extractor( + backbone: resnet.ResNet, + spatial_dims: int, + pretrained_backbone: bool = False, + returned_layers: Sequence[int] = (1, 2, 3), + trainable_backbone_layers: Union[int, None] = None, +): + """ + Constructs a feature extractor network with a ResNet-FPN backbone, used as feature_extractor in RetinaNet. + + Reference: `"Focal Loss for Dense Object Detection" `_. + + The returned feature_extractor network takes an image tensor as inputs, + and outputs a dictionary that maps string to the extracted feature maps (Tensor). + + The input to the returned feature_extractor is expected to be a list of tensors, + each of shape ``[C, H, W]`` or ``[C, H, W, D]``, + one for each image. Different images can have different sizes. + + + Args: + backbone: a ResNet model, used as backbone. + spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images. + pretrained_backbone: whether the backbone has been pre-trained. + returned_layers: returned layers to extract feature maps. Each returned layer should be in the range [1,4]. + len(returned_layers)+1 will be the number of extracted feature maps. + There is an extra maxpooling layer LastLevelMaxPool() appended. + trainable_backbone_layers: number of trainable (not frozen) resnet layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. + When pretrained_backbone is False, this value is set to be 5. + When pretrained_backbone is True, if ``None`` is passed (the default) this value is set to 3. + + Example: + + .. code-block:: python + + from monai.networks.nets import resnet + spatial_dims = 3 # 3D network + backbone = resnet.ResNet( + spatial_dims = spatial_dims, + block = resnet.ResNetBottleneck, + layers = [3, 4, 6, 3], + block_inplanes = resnet.get_inplanes(), + n_input_channels= 1, + conv1_t_stride = (2,2,1), + conv1_t_size = (7,7,7), + ) + # This feature_extractor outputs 4-level feature maps. + # number of output feature maps is len(returned_layers)+1 + feature_extractor = resnet_fpn_feature_extractor( + backbone = backbone, + spatial_dims = spatial_dims, + pretrained_backbone = False, + trainable_backbone_layers = None, + returned_layers = [1,2,3], + ) + model = RetinaNet( + spatial_dims = spatial_dims, + num_classes = 5, + num_anchors = 6, + feature_extractor=feature_extractor, + size_divisible = 32, + ).to(device) + """ + # If pretrained_backbone is False, valid_trainable_backbone_layers = 5. + # If pretrained_backbone is True, valid_trainable_backbone_layers = trainable_backbone_layers or 3 if None. + valid_trainable_backbone_layers: int = _validate_trainable_layers( + pretrained_backbone, trainable_backbone_layers, max_value=5, default_value=3 + ) + + feature_extractor = _resnet_fpn_extractor( + backbone, + spatial_dims, + valid_trainable_backbone_layers, + returned_layers=list(returned_layers), + extra_blocks=None, + ) + return feature_extractor diff --git a/monai/apps/detection/transforms/__init__.py b/monai/apps/detection/transforms/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/transforms/__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/monai/apps/detection/transforms/array.py b/monai/apps/detection/transforms/array.py new file mode 100644 index 0000000000..d5d61f6e43 --- /dev/null +++ b/monai/apps/detection/transforms/array.py @@ -0,0 +1,539 @@ +# 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. +""" +A collection of "vanilla" transforms for box operations +https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design +""" + +from typing import Callable, Optional, Sequence, Tuple, Type, Union + +import numpy as np +import torch + +from monai.config.type_definitions import NdarrayOrTensor +from monai.data.box_utils import ( + BoxMode, + clip_boxes_to_image, + convert_box_mode, + convert_box_to_standard_mode, + get_spatial_dims, + spatial_crop_boxes, +) +from monai.transforms import Rotate90, SpatialCrop +from monai.transforms.transform import Transform +from monai.utils import ensure_tuple, ensure_tuple_rep, fall_back_tuple, look_up_option +from monai.utils.enums import TransformBackends + +from .box_ops import ( + apply_affine_to_boxes, + convert_box_to_mask, + convert_mask_to_box, + flip_boxes, + resize_boxes, + rot90_boxes, + select_labels, + zoom_boxes, +) + +__all__ = [ + "ConvertBoxToStandardMode", + "ConvertBoxMode", + "AffineBox", + "ZoomBox", + "ResizeBox", + "FlipBox", + "ClipBoxToImage", + "BoxToMask", + "MaskToBox", + "SpatialCropBox", + "RotateBox90", +] + + +class ConvertBoxMode(Transform): + """ + This transform converts the boxes in src_mode to the dst_mode. + + Args: + src_mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + dst_mode: target box mode. If it is not given, this func will assume it is ``StandardMode()``. + + Note: + ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, + also represented as "xyxy" for 2D and "xyzxyz" for 3D. + + src_mode and dst_mode can be: + #. str: choose from :class:`~monai.utils.enums.BoxModeName`, for example, + - "xyxy": boxes has format [xmin, ymin, xmax, ymax] + - "xyzxyz": boxes has format [xmin, ymin, zmin, xmax, ymax, zmax] + - "xxyy": boxes has format [xmin, xmax, ymin, ymax] + - "xxyyzz": boxes has format [xmin, xmax, ymin, ymax, zmin, zmax] + - "xyxyzz": boxes has format [xmin, ymin, xmax, ymax, zmin, zmax] + - "xywh": boxes has format [xmin, ymin, xsize, ysize] + - "xyzwhd": boxes has format [xmin, ymin, zmin, xsize, ysize, zsize] + - "ccwh": boxes has format [xcenter, ycenter, xsize, ysize] + - "cccwhd": boxes has format [xcenter, ycenter, zcenter, xsize, ysize, zsize] + #. BoxMode class: choose from the subclasses of :class:`~monai.data.box_utils.BoxMode`, for example, + - CornerCornerModeTypeA: equivalent to "xyxy" or "xyzxyz" + - CornerCornerModeTypeB: equivalent to "xxyy" or "xxyyzz" + - CornerCornerModeTypeC: equivalent to "xyxy" or "xyxyzz" + - CornerSizeMode: equivalent to "xywh" or "xyzwhd" + - CenterSizeMode: equivalent to "ccwh" or "cccwhd" + #. BoxMode object: choose from the subclasses of :class:`~monai.data.box_utils.BoxMode`, for example, + - CornerCornerModeTypeA(): equivalent to "xyxy" or "xyzxyz" + - CornerCornerModeTypeB(): equivalent to "xxyy" or "xxyyzz" + - CornerCornerModeTypeC(): equivalent to "xyxy" or "xyxyzz" + - CornerSizeMode(): equivalent to "xywh" or "xyzwhd" + - CenterSizeMode(): equivalent to "ccwh" or "cccwhd" + #. None: will assume mode is ``StandardMode()`` + + Example: + .. code-block:: python + + boxes = torch.ones(10,4) + # convert boxes with format [xmin, ymin, xmax, ymax] to [xcenter, ycenter, xsize, ysize]. + box_converter = ConvertBoxMode(src_mode="xyxy", dst_mode="ccwh") + box_converter(boxes) + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + src_mode: Union[str, BoxMode, Type[BoxMode], None] = None, + dst_mode: Union[str, BoxMode, Type[BoxMode], None] = None, + ) -> None: + self.src_mode = src_mode + self.dst_mode = dst_mode + + def __call__(self, boxes: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Converts the boxes in src_mode to the dst_mode. + + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + bounding boxes with target mode, with same data type as ``boxes``, does not share memory with ``boxes`` + """ + return convert_box_mode(boxes, src_mode=self.src_mode, dst_mode=self.dst_mode) + + +class ConvertBoxToStandardMode(Transform): + """ + Convert given boxes to standard mode. + Standard mode is "xyxy" or "xyzxyz", + representing box format of [xmin, ymin, xmax, ymax] or [xmin, ymin, zmin, xmax, ymax, zmax]. + + Args: + mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``src_mode`` in :class:`~monai.apps.detection.transforms.array.ConvertBoxMode` . + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + # convert boxes with format [xmin, xmax, ymin, ymax, zmin, zmax] to [xmin, ymin, zmin, xmax, ymax, zmax] + box_converter = ConvertBoxToStandardMode(mode="xxyyzz") + box_converter(boxes) + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, mode: Union[str, BoxMode, Type[BoxMode], None] = None) -> None: + self.mode = mode + + def __call__(self, boxes: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Convert given boxes to standard mode. + Standard mode is "xyxy" or "xyzxyz", + representing box format of [xmin, ymin, xmax, ymax] or [xmin, ymin, zmin, xmax, ymax, zmax]. + + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + bounding boxes with standard mode, with same data type as ``boxes``, does not share memory with ``boxes`` + """ + return convert_box_to_standard_mode(boxes, mode=self.mode) + + +class AffineBox(Transform): + """ + Applies affine matrix to the boxes + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __call__(self, boxes: NdarrayOrTensor, affine: Union[NdarrayOrTensor, None]) -> NdarrayOrTensor: # type: ignore + """ + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + affine: affine matrix to be applied to the box coordinate + """ + if affine is None: + return boxes + + return apply_affine_to_boxes(boxes, affine=affine) + + +class ZoomBox(Transform): + """ + Zooms an ND Box with same padding or slicing setting with Zoom(). + + Args: + zoom: The zoom factor along the spatial axes. + If a float, zoom is the same for each spatial axis. + If a sequence, zoom should contain one value for each spatial axis. + keep_size: Should keep original size (padding/slicing if needed), default is True. + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, zoom: Union[Sequence[float], float], keep_size: bool = False, **kwargs) -> None: + self.zoom = zoom + self.keep_size = keep_size + self.kwargs = kwargs + + def __call__(self, boxes: torch.Tensor, src_spatial_size: Union[Sequence[int], int, None] = None): + """ + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + src_spatial_size: original image spatial size before zooming, used only when keep_size=True. + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + self._zoom = ensure_tuple_rep(self.zoom, spatial_dims) # match the spatial image dim + + if not self.keep_size: + return zoom_boxes(boxes, self._zoom) + + if src_spatial_size is None: + raise ValueError("keep_size=True, src_spatial_size must be provided.") + + src_spatial_size = ensure_tuple_rep(src_spatial_size, spatial_dims) + dst_spatial_size = [int(round(z * ss)) for z, ss in zip(self._zoom, src_spatial_size)] + self._zoom = tuple(ds / float(ss) for ss, ds in zip(src_spatial_size, dst_spatial_size)) + zoomed_boxes = zoom_boxes(boxes, self._zoom) + + # See also keep_size in monai.transforms.spatial.array.Zoom() + if not np.allclose(np.array(src_spatial_size), np.array(dst_spatial_size)): + for axis, (od, zd) in enumerate(zip(src_spatial_size, dst_spatial_size)): + diff = od - zd + half = abs(diff) // 2 + if diff > 0: # need padding (half, diff - half) + zoomed_boxes[:, axis] = zoomed_boxes[:, axis] + half + zoomed_boxes[:, axis + spatial_dims] = zoomed_boxes[:, axis + spatial_dims] + half + elif diff < 0: # need slicing (half, half + od) + zoomed_boxes[:, axis] = zoomed_boxes[:, axis] - half + zoomed_boxes[:, axis + spatial_dims] = zoomed_boxes[:, axis + spatial_dims] - half + return zoomed_boxes + + +class ResizeBox(Transform): + """ + Resize the input boxes when the corresponding image is + resized to given spatial size (with scaling, not cropping/padding). + + Args: + spatial_size: expected shape of spatial dimensions after resize operation. + if some components of the `spatial_size` are non-positive values, the transform will use the + corresponding components of img size. For example, `spatial_size=(32, -1)` will be adapted + to `(32, 64)` if the second spatial dimension size of img is `64`. + size_mode: should be "all" or "longest", if "all", will use `spatial_size` for all the spatial dims, + if "longest", rescale the image so that only the longest side is equal to specified `spatial_size`, + which must be an int number in this case, keeping the aspect ratio of the initial image, refer to: + https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/ + #albumentations.augmentations.geometric.resize.LongestMaxSize. + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, spatial_size: Union[Sequence[int], int], size_mode: str = "all", **kwargs) -> None: + self.size_mode = look_up_option(size_mode, ["all", "longest"]) + self.spatial_size = spatial_size + + def __call__(self, boxes: NdarrayOrTensor, src_spatial_size: Union[Sequence[int], int]): # type: ignore + """ + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + src_spatial_size: original image spatial size before resizing. + + Raises: + ValueError: When ``self.spatial_size`` length is less than ``boxes`` spatial dimensions. + """ + input_ndim = get_spatial_dims(boxes=boxes) # spatial ndim + src_spatial_size_ = ensure_tuple_rep(src_spatial_size, input_ndim) + + if self.size_mode == "all": + # spatial_size must be a Sequence if size_mode is 'all' + output_ndim = len(ensure_tuple(self.spatial_size)) + if output_ndim != input_ndim: + raise ValueError( + "len(spatial_size) must be greater or equal to img spatial dimensions, " + f"got spatial_size={output_ndim} img={input_ndim}." + ) + spatial_size_ = fall_back_tuple(self.spatial_size, src_spatial_size_) + else: # for the "longest" mode + if not isinstance(self.spatial_size, int): + raise ValueError("spatial_size must be an int number if size_mode is 'longest'.") + scale = self.spatial_size / max(src_spatial_size_) + spatial_size_ = tuple(int(round(s * scale)) for s in src_spatial_size_) + + return resize_boxes(boxes, src_spatial_size_, spatial_size_) + + +class FlipBox(Transform): + """ + Reverses the box coordinates along the given spatial axis. Preserves shape. + + Args: + spatial_axis: spatial axes along which to flip over. Default is None. + The default `axis=None` will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, spatial_axis: Optional[Union[Sequence[int], int]] = None) -> None: + self.spatial_axis = spatial_axis + + def __call__(self, boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], int]): # type: ignore + """ + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: image spatial size. + """ + + return flip_boxes(boxes, spatial_size=spatial_size, flip_axes=self.spatial_axis) + + +class ClipBoxToImage(Transform): + """ + Clip the bounding boxes and the associated labels/scores to make sure they are within the image. + There might be multiple arrays of labels/scores associated with one array of boxes. + + Args: + remove_empty: whether to remove the boxes and corresponding labels that are actually empty + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, remove_empty: bool = False) -> None: + self.remove_empty = remove_empty + + def __call__( # type: ignore + self, + boxes: NdarrayOrTensor, + labels: Union[Sequence[NdarrayOrTensor], NdarrayOrTensor], + spatial_size: Union[Sequence[int], int], + ) -> Tuple[NdarrayOrTensor, Union[Tuple, NdarrayOrTensor]]: + """ + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + labels: Sequence of array. Each element represents classification labels or scores + corresponding to ``boxes``, sized (N,). + spatial_size: The spatial size of the image where the boxes are attached. len(spatial_size) should be in [2, 3]. + + Returns: + - clipped boxes, does not share memory with original boxes + - clipped labels, does not share memory with original labels + + Example: + .. code-block:: python + + box_clipper = ClipBoxToImage(remove_empty=True) + boxes = torch.ones(2, 6) + class_labels = torch.Tensor([0, 1]) + pred_scores = torch.Tensor([[0.4,0.3,0.3], [0.5,0.1,0.4]]) + labels = (class_labels, pred_scores) + spatial_size = [32, 32, 32] + boxes_clip, labels_clip_tuple = box_clipper(boxes, labels, spatial_size) + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + spatial_size = ensure_tuple_rep(spatial_size, spatial_dims) # match the spatial image dim + + boxes_clip, keep = clip_boxes_to_image(boxes, spatial_size, self.remove_empty) + return boxes_clip, select_labels(labels, keep) + + +class BoxToMask(Transform): + """ + Convert box to int16 mask image, which has the same size with the input image. + + Args: + bg_label: background labels for the output mask image, make sure it is smaller than any foreground(fg) labels. + ellipse_mask: bool. + + - If True, it assumes the object shape is close to ellipse or ellipsoid. + - If False, it assumes the object shape is close to rectangle or cube and well occupies the bounding box. + - If the users are going to apply random rotation as data augmentation, we suggest setting ellipse_mask=True + See also Kalra et al. "Towards Rotation Invariance in Object Detection", ICCV 2021. + """ + + backend = [TransformBackends.NUMPY] + + def __init__(self, bg_label: int = -1, ellipse_mask: bool = False) -> None: + self.bg_label = bg_label + self.ellipse_mask = ellipse_mask + + def __call__( # type: ignore + self, boxes: NdarrayOrTensor, labels: NdarrayOrTensor, spatial_size: Union[Sequence[int], int] + ) -> NdarrayOrTensor: + """ + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. + labels: classification foreground(fg) labels corresponding to `boxes`, dtype should be int, sized (N,). + spatial_size: image spatial size. + + Return: + - int16 array, sized (num_box, H, W). Each channel represents a box. + The foreground region in channel c has intensity of labels[c]. + The background intensity is bg_label. + """ + return convert_box_to_mask(boxes, labels, spatial_size, self.bg_label, self.ellipse_mask) + + +class MaskToBox(Transform): + """ + Convert int16 mask image to box, which has the same size with the input image. + Pairs with :py:class:`monai.apps.detection.transforms.array.BoxToMask`. + Please make sure the same ``min_fg_label`` is used when using the two transforms in pairs. + + Args: + bg_label: background labels for the output mask image, make sure it is smaller than any foreground(fg) labels. + box_dtype: output dtype for boxes + label_dtype: output dtype for labels + """ + + backend = [TransformBackends.NUMPY] + + def __init__(self, bg_label: int = -1, box_dtype=torch.float32, label_dtype=torch.long) -> None: + self.bg_label = bg_label + self.box_dtype = box_dtype + self.label_dtype = label_dtype + + def __call__(self, boxes_mask: NdarrayOrTensor) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + """ + Args: + boxes_mask: int16 array, sized (num_box, H, W). Each channel represents a box. + The foreground region in channel c has intensity of labels[c]. + The background intensity is bg_label. + + Return: + - bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. + - classification foreground(fg) labels, dtype should be int, sized (N,). + """ + return convert_mask_to_box(boxes_mask, self.bg_label, self.box_dtype, self.label_dtype) + + +class SpatialCropBox(SpatialCrop): + """ + General purpose box cropper when the corresponding image is cropped by SpatialCrop(*) with the same ROI. + The difference is that we do not support negative indexing for roi_slices. + + If a dimension of the expected ROI size is bigger than the input image size, will not crop that dimension. + So the cropped result may be smaller than the expected ROI, and the cropped results of several images may + not have exactly the same shape. + It can support to crop ND spatial boxes. + + The cropped region can be parameterised in various ways: + - a list of slices for each spatial dimension (do not allow for use of negative indexing) + - a spatial center and size + - the start and end coordinates of the ROI + + Args: + roi_center: voxel coordinates for center of the crop ROI. + roi_size: size of the crop ROI, if a dimension of ROI size is bigger than image size, + will not crop that dimension of the image. + roi_start: voxel coordinates for start of the crop ROI. + roi_end: voxel coordinates for end of the crop ROI, if a coordinate is out of image, + use the end coordinate of image. + roi_slices: list of slices for each of the spatial dimensions. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + roi_center: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_size: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_start: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_end: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_slices: Optional[Sequence[slice]] = None, + ) -> None: + super().__init__(roi_center, roi_size, roi_start, roi_end, roi_slices) + for s in self.slices: + if s.start < 0 or s.stop < 0 or (s.step is not None and s.step < 0): + raise ValueError("Currently negative indexing is not supported for SpatialCropBox.") + + def __call__( # type: ignore + self, boxes: NdarrayOrTensor, labels: Union[Sequence[NdarrayOrTensor], NdarrayOrTensor] + ): + """ + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + labels: Sequence of array. Each element represents classification labels or scores + + Returns: + - cropped boxes, does not share memory with original boxes + - cropped labels, does not share memory with original labels + + Example: + .. code-block:: python + + box_cropper = SpatialCropPadBox(roi_start=[0, 1, 4], roi_end=[21, 15, 8]) + boxes = torch.ones(2, 6) + class_labels = torch.Tensor([0, 1]) + pred_scores = torch.Tensor([[0.4,0.3,0.3], [0.5,0.1,0.4]]) + labels = (class_labels, pred_scores) + boxes_crop, labels_crop_tuple = box_cropper(boxes, labels) + """ + spatial_dims = min(len(self.slices), get_spatial_dims(boxes=boxes)) # spatial dims + boxes_crop, keep = spatial_crop_boxes( + boxes, + [self.slices[axis].start for axis in range(spatial_dims)], + [self.slices[axis].stop for axis in range(spatial_dims)], + ) + return boxes_crop, select_labels(labels, keep) + + +class RotateBox90(Rotate90): + """ + Rotate a boxes by 90 degrees in the plane specified by `axes`. + See box_ops.rot90_boxes for additional details + + Args: + k: number of times to rotate by 90 degrees. + spatial_axes: 2 int numbers, defines the plane to rotate with 2 spatial axes. + Default: (0, 1), this is the first two axis in spatial dimensions. + If axis is negative it counts from the last to the first axis. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, k: int = 1, spatial_axes: Tuple[int, int] = (0, 1)) -> None: + super().__init__(k, spatial_axes) + + def __call__(self, boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], int]): # type: ignore + """ + Args: + img: channel first array, must have shape: (num_channels, H[, W, ..., ]), + """ + rot90: Callable = rot90_boxes + out: NdarrayOrTensor = rot90(boxes, spatial_size, self.k, self.spatial_axes) + return out diff --git a/monai/apps/detection/transforms/box_ops.py b/monai/apps/detection/transforms/box_ops.py new file mode 100644 index 0000000000..1562f28b29 --- /dev/null +++ b/monai/apps/detection/transforms/box_ops.py @@ -0,0 +1,425 @@ +# 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 copy import deepcopy +from typing import Optional, Sequence, Tuple, Union + +import numpy as np +import torch + +from monai.config.type_definitions import NdarrayOrTensor +from monai.data.box_utils import COMPUTE_DTYPE, TO_REMOVE, get_spatial_dims +from monai.transforms import Resize +from monai.transforms.utils import create_scale +from monai.utils import look_up_option, optional_import +from monai.utils.misc import ensure_tuple, ensure_tuple_rep +from monai.utils.type_conversion import convert_data_type, convert_to_dst_type + +scipy, _ = optional_import("scipy") + + +def _apply_affine_to_points(points: torch.Tensor, affine: torch.Tensor, include_shift: bool = True) -> torch.Tensor: + """ + This internal function applies affine matrices to the point coordinate + + Args: + points: point coordinates, Nx2 or Nx3 torch tensor or ndarray, representing [x, y] or [x, y, z] + affine: affine matrix to be applied to the point coordinates, sized (spatial_dims+1,spatial_dims+1) + include_shift: default True, whether the function apply translation (shift) in the affine transform + + Returns: + transformed point coordinates, with same data type as ``points``, does not share memory with ``points`` + """ + + spatial_dims = get_spatial_dims(points=points) + + # compute new points + if include_shift: + # append 1 to form Nx(spatial_dims+1) vector, then transpose + points_affine = torch.cat( + [points, torch.ones(points.shape[0], 1, device=points.device, dtype=points.dtype)], dim=1 + ).transpose(0, 1) + # apply affine + points_affine = torch.matmul(affine, points_affine) + # remove appended 1 and transpose back + points_affine = points_affine[:spatial_dims, :].transpose(0, 1) + else: + points_affine = points.transpose(0, 1) + points_affine = torch.matmul(affine[:spatial_dims, :spatial_dims], points_affine) + points_affine = points_affine.transpose(0, 1) + + return points_affine + + +def apply_affine_to_boxes(boxes: NdarrayOrTensor, affine: NdarrayOrTensor) -> NdarrayOrTensor: + """ + This function applies affine matrices to the boxes + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be StandardMode + affine: affine matrix to be applied to the box coordinates, sized (spatial_dims+1,spatial_dims+1) + + Returns: + returned affine transformed boxes, with same data type as ``boxes``, does not share memory with ``boxes`` + """ + + # convert numpy to tensor if needed + boxes_t, *_ = convert_data_type(boxes, torch.Tensor) + + # some operation does not support torch.float16 + # convert to float32 + + boxes_t = boxes_t.to(dtype=COMPUTE_DTYPE) + affine_t, *_ = convert_to_dst_type(src=affine, dst=boxes_t) + + spatial_dims = get_spatial_dims(boxes=boxes_t) + + # affine transform left top and bottom right points + # might flipped, thus lt may not be left top any more + lt: torch.Tensor = _apply_affine_to_points(boxes_t[:, :spatial_dims], affine_t, include_shift=True) + rb: torch.Tensor = _apply_affine_to_points(boxes_t[:, spatial_dims:], affine_t, include_shift=True) + + # make sure lt_new is left top, and rb_new is bottom right + lt_new, _ = torch.min(torch.stack([lt, rb], dim=2), dim=2) + rb_new, _ = torch.max(torch.stack([lt, rb], dim=2), dim=2) + + boxes_t_affine = torch.cat([lt_new, rb_new], dim=1) + + # convert tensor back to numpy if needed + boxes_affine: NdarrayOrTensor + boxes_affine, *_ = convert_to_dst_type(src=boxes_t_affine, dst=boxes) + return boxes_affine + + +def zoom_boxes(boxes: NdarrayOrTensor, zoom: Union[Sequence[float], float]): + """ + Zoom boxes + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be StandardMode + zoom: The zoom factor along the spatial axes. + If a float, zoom is the same for each spatial axis. + If a sequence, zoom should contain one value for each spatial axis. + + Returns: + zoomed boxes, with same data type as ``boxes``, does not share memory with ``boxes`` + + Example: + .. code-block:: python + + boxes = torch.ones(1,4) + zoom_boxes(boxes, zoom=[0.5,2.2]) # will return tensor([[0.5, 2.2, 0.5, 2.2]]) + """ + spatial_dims = get_spatial_dims(boxes=boxes) + + # generate affine transform corresponding to ``zoom`` + affine = create_scale(spatial_dims=spatial_dims, scaling_factor=zoom) + + return apply_affine_to_boxes(boxes=boxes, affine=affine) + + +def resize_boxes( + boxes: NdarrayOrTensor, src_spatial_size: Union[Sequence[int], int], dst_spatial_size: Union[Sequence[int], int] +): + """ + Resize boxes when the corresponding image is resized + + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + src_spatial_size: source image spatial size. + dst_spatial_size: target image spatial size. + + Returns: + resized boxes, with same data type as ``boxes``, does not share memory with ``boxes`` + + Example: + .. code-block:: python + + boxes = torch.ones(1,4) + src_spatial_size = [100, 100] + dst_spatial_size = [128, 256] + resize_boxes(boxes, src_spatial_size, dst_spatial_size) # will return tensor([[1.28, 2.56, 1.28, 2.56]]) + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + + src_spatial_size = ensure_tuple_rep(src_spatial_size, spatial_dims) + dst_spatial_size = ensure_tuple_rep(dst_spatial_size, spatial_dims) + + zoom = [dst_spatial_size[axis] / float(src_spatial_size[axis]) for axis in range(spatial_dims)] + + return zoom_boxes(boxes=boxes, zoom=zoom) + + +def flip_boxes( + boxes: NdarrayOrTensor, + spatial_size: Union[Sequence[int], int], + flip_axes: Optional[Union[Sequence[int], int]] = None, +) -> NdarrayOrTensor: + """ + Flip boxes when the corresponding image is flipped + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: image spatial size. + flip_axes: spatial axes along which to flip over. Default is None. + The default `axis=None` will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + + Returns: + flipped boxes, with same data type as ``boxes``, does not share memory with ``boxes`` + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + spatial_size = ensure_tuple_rep(spatial_size, spatial_dims) + if flip_axes is None: + flip_axes = tuple(range(0, spatial_dims)) + flip_axes = ensure_tuple(flip_axes) + + # flip box + _flip_boxes = deepcopy(boxes) + for axis in flip_axes: + _flip_boxes[:, axis + spatial_dims] = spatial_size[axis] - boxes[:, axis] - TO_REMOVE + _flip_boxes[:, axis] = spatial_size[axis] - boxes[:, axis + spatial_dims] - TO_REMOVE + + return _flip_boxes + + +def convert_box_to_mask( + boxes: NdarrayOrTensor, + labels: NdarrayOrTensor, + spatial_size: Union[Sequence[int], int], + bg_label: int = -1, + ellipse_mask: bool = False, +) -> NdarrayOrTensor: + """ + Convert box to int16 mask image, which has the same size with the input image. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. + labels: classification foreground(fg) labels corresponding to `boxes`, dtype should be int, sized (N,). + spatial_size: image spatial size. + bg_label: background labels for the output mask image, make sure it is smaller than any fg labels. + ellipse_mask: bool. + + - If True, it assumes the object shape is close to ellipse or ellipsoid. + - If False, it assumes the object shape is close to rectangle or cube and well occupies the bounding box. + - If the users are going to apply random rotation as data augmentation, we suggest setting ellipse_mask=True + See also Kalra et al. "Towards Rotation Invariance in Object Detection", ICCV 2021. + + Return: + - int16 array, sized (num_box, H, W). Each channel represents a box. + The foreground region in channel c has intensity of labels[c]. + The background intensity is bg_label. + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + spatial_size = ensure_tuple_rep(spatial_size, spatial_dims) + + # if no box, return empty mask + if len(labels) == 0: + boxes_mask_np = np.ones((1,) + spatial_size, dtype=np.int16) * np.int16(bg_label) + boxes_mask, *_ = convert_to_dst_type(src=boxes_mask_np, dst=boxes, dtype=torch.int16) + return boxes_mask + + # bg_label should be smaller than labels + if bg_label >= min(labels): + raise ValueError( + f"bg_label should be smaller than any foreground box labels.\n" + f"min(labels)={min(labels)}, while bg_label={bg_label}" + ) + + if labels.shape[0] != boxes.shape[0]: + raise ValueError("Number of labels should equal to number of boxes.") + + # allocate memory for boxes_mask_np + boxes_mask_np = np.ones((labels.shape[0],) + spatial_size, dtype=np.int16) * np.int16(bg_label) + + boxes_np: np.ndarray = convert_data_type(boxes, np.ndarray, dtype=np.int32)[0] + labels_np, *_ = convert_to_dst_type(src=labels, dst=boxes_np) + for b in range(boxes_np.shape[0]): + # generate a foreground mask + box_size = [boxes_np[b, axis + spatial_dims] - boxes_np[b, axis] for axis in range(spatial_dims)] + if ellipse_mask: + # initialize a square/cube mask + max_box_size = max(box_size) + radius = max_box_size / 2.0 + center = (max_box_size - 1) / 2.0 + boxes_only_mask = np.ones([max_box_size] * spatial_dims, dtype=np.int16) * np.int16(bg_label) + # apply label intensity to circle/ball foreground + ranges = tuple(slice(0, max_box_size) for _ in range(spatial_dims)) + dist_from_center = sum((grid - center) ** 2 for grid in np.ogrid[ranges]) + boxes_only_mask[dist_from_center <= radius**2] = np.int16(labels_np[b]) + # squeeze it to a ellipse/ellipsoid mask + resizer = Resize(spatial_size=box_size, mode="nearest", anti_aliasing=False) + boxes_only_mask = resizer(boxes_only_mask[None])[0] # type: ignore + else: + # generate a rect mask + boxes_only_mask = np.ones(box_size, dtype=np.int16) * np.int16(labels_np[b]) + # apply to global mask + slicing = [b] + slicing.extend(slice(boxes_np[b, d], boxes_np[b, d + spatial_dims]) for d in range(spatial_dims)) # type:ignore + boxes_mask_np[tuple(slicing)] = boxes_only_mask + return convert_to_dst_type(src=boxes_mask_np, dst=boxes, dtype=torch.int16)[0] + + +def convert_mask_to_box( + boxes_mask: NdarrayOrTensor, bg_label: int = -1, box_dtype=torch.float32, label_dtype=torch.long +) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + """ + Convert int16 mask image to box, which has the same size with the input image + + Args: + boxes_mask: int16 array, sized (num_box, H, W). Each channel represents a box. + The foreground region in channel c has intensity of labels[c]. + The background intensity is bg_label. + bg_label: background labels for the boxes_mask + box_dtype: output dtype for boxes + label_dtype: output dtype for labels + + Return: + - bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. + - classification foreground(fg) labels, dtype should be int, sized (N,). + """ + look_up_option(len(boxes_mask.shape), [3, 4]) + spatial_size = list(boxes_mask.shape[1:]) + spatial_dims = get_spatial_dims(spatial_size=spatial_size) + + boxes_mask_np, *_ = convert_data_type(boxes_mask, np.ndarray) + + boxes_list = [] + labels_list = [] + for b in range(boxes_mask_np.shape[0]): + fg_indices = np.nonzero(boxes_mask_np[b, ...] - bg_label) + if fg_indices[0].shape[0] == 0: + continue + boxes_b = [] + for fd_i in fg_indices: + boxes_b.append(min(fd_i)) # top left corner + for fd_i in fg_indices: + boxes_b.append(max(fd_i) + 1 - TO_REMOVE) # bottom right corner + boxes_list.append(boxes_b) + if spatial_dims == 2: + labels_list.append(boxes_mask_np[b, fg_indices[0][0], fg_indices[1][0]]) + if spatial_dims == 3: + labels_list.append(boxes_mask_np[b, fg_indices[0][0], fg_indices[1][0], fg_indices[2][0]]) + + if len(boxes_list) == 0: + boxes_np, labels_np = np.zeros([0, 2 * spatial_dims]), np.zeros([0]) + else: + boxes_np, labels_np = np.asarray(boxes_list), np.asarray(labels_list) + boxes, *_ = convert_to_dst_type(src=boxes_np, dst=boxes_mask, dtype=box_dtype) + labels, *_ = convert_to_dst_type(src=labels_np, dst=boxes_mask, dtype=label_dtype) + return boxes, labels + + +def select_labels( + labels: Union[Sequence[NdarrayOrTensor], NdarrayOrTensor], keep: NdarrayOrTensor +) -> Union[Tuple, NdarrayOrTensor]: + """ + For element in labels, select indice keep from it. + + Args: + labels: Sequence of array. Each element represents classification labels or scores + corresponding to ``boxes``, sized (N,). + keep: the indices to keep, same length with each element in labels. + + Return: + selected labels, does not share memory with original labels. + """ + labels_tuple = ensure_tuple(labels, True) + + labels_select_list = [] + keep_t: torch.Tensor = convert_data_type(keep, torch.Tensor)[0] + for i in range(len(labels_tuple)): + labels_t: torch.Tensor = convert_data_type(labels_tuple[i], torch.Tensor)[0] + labels_t = labels_t[keep_t, ...] + labels_select_list.append(convert_to_dst_type(src=labels_t, dst=labels_tuple[i])[0]) + + if isinstance(labels, (torch.Tensor, np.ndarray)): + return labels_select_list[0] # type: ignore + + return tuple(labels_select_list) + + +def swapaxes_boxes(boxes: NdarrayOrTensor, axis1: int, axis2: int) -> NdarrayOrTensor: + """ + Interchange two axes of boxes. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + axis1: First axis. + axis2: Second axis. + + Returns: + boxes with two axes interchanged. + + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + boxes_swap: NdarrayOrTensor = deepcopy(boxes) + boxes_swap[:, [axis1, axis2]] = boxes_swap[:, [axis2, axis1]] # type: ignore + boxes_swap[:, [spatial_dims + axis1, spatial_dims + axis2]] = boxes_swap[ # type: ignore + :, [spatial_dims + axis2, spatial_dims + axis1] + ] + return boxes_swap + + +def rot90_boxes( + boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], int], k: int = 1, axes: Tuple[int, int] = (0, 1) +) -> NdarrayOrTensor: + """ + Rotate boxes by 90 degrees in the plane specified by axes. + Rotation direction is from the first towards the second axis. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: image spatial size. + k : number of times the array is rotated by 90 degrees. + axes: (2,) array_like + The array is rotated in the plane defined by the axes. Axes must be different. + + Returns: + A rotated view of `boxes`. + + Notes: + ``rot90_boxes(boxes, spatial_size, k=1, axes=(1,0))`` is the reverse of + ``rot90_boxes(boxes, spatial_size, k=1, axes=(0,1))`` + ``rot90_boxes(boxes, spatial_size, k=1, axes=(1,0))`` is equivalent to + ``rot90_boxes(boxes, spatial_size, k=-1, axes=(0,1))`` + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + spatial_size_ = list(ensure_tuple_rep(spatial_size, spatial_dims)) + + axes = ensure_tuple(axes) # type: ignore + + if len(axes) != 2: + raise ValueError("len(axes) must be 2.") + + if axes[0] == axes[1] or abs(axes[0] - axes[1]) == spatial_dims: + raise ValueError("Axes must be different.") + + if axes[0] >= spatial_dims or axes[0] < -spatial_dims or axes[1] >= spatial_dims or axes[1] < -spatial_dims: + raise ValueError(f"Axes={axes} out of range for array of ndim={spatial_dims}.") + + k %= 4 + + if k == 0: + return boxes + if k == 2: + return flip_boxes(flip_boxes(boxes, spatial_size_, axes[0]), spatial_size_, axes[1]) + + if k == 1: + boxes_ = flip_boxes(boxes, spatial_size_, axes[1]) + return swapaxes_boxes(boxes_, axes[0], axes[1]) + else: + # k == 3 + boxes_ = swapaxes_boxes(boxes, axes[0], axes[1]) + spatial_size_[axes[0]], spatial_size_[axes[1]] = spatial_size_[axes[1]], spatial_size_[axes[0]] + return flip_boxes(boxes_, spatial_size_, axes[1]) diff --git a/monai/apps/detection/transforms/dictionary.py b/monai/apps/detection/transforms/dictionary.py new file mode 100644 index 0000000000..cb08d0ed69 --- /dev/null +++ b/monai/apps/detection/transforms/dictionary.py @@ -0,0 +1,1398 @@ +# 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. +""" +A collection of dictionary-based wrappers around the "vanilla" transforms for box operations +defined in :py:class:`monai.apps.detection.transforms.array`. + +Class names are ended with 'd' to denote dictionary-based transforms. +""" +from copy import deepcopy +from enum import Enum +from typing import Dict, Hashable, List, Mapping, Optional, Sequence, Tuple, Type, Union + +import numpy as np +import torch + +from monai.apps.detection.transforms.array import ( + AffineBox, + BoxToMask, + ClipBoxToImage, + ConvertBoxMode, + ConvertBoxToStandardMode, + FlipBox, + MaskToBox, + RotateBox90, + SpatialCropBox, + ZoomBox, +) +from monai.apps.detection.transforms.box_ops import convert_box_to_mask +from monai.config import KeysCollection, SequenceStr +from monai.config.type_definitions import NdarrayOrTensor +from monai.data.box_utils import COMPUTE_DTYPE, BoxMode, clip_boxes_to_image +from monai.data.utils import orientation_ras_lps +from monai.transforms import Flip, RandFlip, RandRotate90d, RandZoom, Rotate90, SpatialCrop, SpatialPad, Zoom +from monai.transforms.inverse import InvertibleTransform +from monai.transforms.transform import MapTransform, Randomizable, RandomizableTransform +from monai.transforms.utils import generate_pos_neg_label_crop_centers, map_binary_to_indices +from monai.utils import ImageMetaKey as Key +from monai.utils import InterpolateMode, NumpyPadMode, ensure_tuple, ensure_tuple_rep +from monai.utils.enums import PostFix, TraceKeys +from monai.utils.type_conversion import convert_data_type + +__all__ = [ + "ConvertBoxModed", + "ConvertBoxModeD", + "ConvertBoxModeDict", + "ConvertBoxToStandardModed", + "ConvertBoxToStandardModeD", + "ConvertBoxToStandardModeDict", + "AffineBoxToImageCoordinated", + "AffineBoxToImageCoordinateD", + "AffineBoxToImageCoordinateDict", + "ZoomBoxd", + "ZoomBoxD", + "ZoomBoxDict", + "RandZoomBoxd", + "RandZoomBoxD", + "RandZoomBoxDict", + "FlipBoxd", + "FlipBoxD", + "FlipBoxDict", + "RandFlipBoxd", + "RandFlipBoxD", + "RandFlipBoxDict", + "ClipBoxToImaged", + "ClipBoxToImageD", + "ClipBoxToImageDict", + "BoxToMaskd", + "BoxToMaskD", + "BoxToMaskDict", + "MaskToBoxd", + "MaskToBoxD", + "MaskToBoxDict", + "RandCropBoxByPosNegLabeld", + "RandCropBoxByPosNegLabelD", + "RandCropBoxByPosNegLabelDict", + "RotateBox90d", + "RotateBox90D", + "RotateBox90Dict", + "RandRotateBox90d", + "RandRotateBox90D", + "RandRotateBox90Dict", +] + +DEFAULT_POST_FIX = PostFix.meta() + + +class ConvertBoxModed(MapTransform, InvertibleTransform): + """ + Dictionary-based wrapper of :py:class:`monai.apps.detection.transforms.array.ConvertBoxMode`. + + This transform converts the boxes in src_mode to the dst_mode. + + Example: + .. code-block:: python + + data = {"boxes": torch.ones(10,4)} + # convert boxes with format [xmin, ymin, xmax, ymax] to [xcenter, ycenter, xsize, ysize]. + box_converter = ConvertBoxModed(box_keys=["boxes"], src_mode="xyxy", dst_mode="ccwh") + box_converter(data) + """ + + def __init__( + self, + box_keys: KeysCollection, + src_mode: Union[str, BoxMode, Type[BoxMode], None] = None, + dst_mode: Union[str, BoxMode, Type[BoxMode], None] = None, + allow_missing_keys: bool = False, + ) -> None: + """ + Args: + box_keys: Keys to pick data for transformation. + src_mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``src_mode`` in :class:`~monai.apps.detection.transforms.array.ConvertBoxMode` . + dst_mode: target box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``src_mode`` in :class:`~monai.apps.detection.transforms.array.ConvertBoxMode` . + allow_missing_keys: don't raise exception if key is missing. + + See also :py:class:`monai.apps.detection,transforms.array.ConvertBoxMode` + """ + super().__init__(box_keys, allow_missing_keys) + self.converter = ConvertBoxMode(src_mode=src_mode, dst_mode=dst_mode) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + self.push_transform(d, key, extra_info={"src": self.converter.src_mode, "dst": self.converter.dst_mode}) + d[key] = self.converter(d[key]) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + for key in self.key_iterator(d): + tr = self.get_most_recent_transform(d, key) + src_mode, dst_mode = tr[TraceKeys.EXTRA_INFO]["src"], tr[TraceKeys.EXTRA_INFO]["dst"] + inverse_converter = ConvertBoxMode(src_mode=dst_mode, dst_mode=src_mode) + # Inverse is same as forward + d[key] = inverse_converter(d[key]) + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class ConvertBoxToStandardModed(MapTransform, InvertibleTransform): + """ + Dictionary-based wrapper of :py:class:`monai.apps.detection.transforms.array.ConvertBoxToStandardMode`. + + Convert given boxes to standard mode. + Standard mode is "xyxy" or "xyzxyz", + representing box format of [xmin, ymin, xmax, ymax] or [xmin, ymin, zmin, xmax, ymax, zmax]. + + Example: + .. code-block:: python + + data = {"boxes": torch.ones(10,6)} + # convert boxes with format [xmin, xmax, ymin, ymax, zmin, zmax] to [xmin, ymin, zmin, xmax, ymax, zmax] + box_converter = ConvertBoxToStandardModed(box_keys=["boxes"], mode="xxyyzz") + box_converter(data) + """ + + def __init__( + self, + box_keys: KeysCollection, + mode: Union[str, BoxMode, Type[BoxMode], None] = None, + allow_missing_keys: bool = False, + ) -> None: + """ + Args: + box_keys: Keys to pick data for transformation. + mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``src_mode`` in :class:`~monai.apps.detection.transforms.array.ConvertBoxMode` . + allow_missing_keys: don't raise exception if key is missing. + + See also :py:class:`monai.apps.detection,transforms.array.ConvertBoxToStandardMode` + """ + super().__init__(box_keys, allow_missing_keys) + self.converter = ConvertBoxToStandardMode(mode=mode) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + self.push_transform(d, key, extra_info={"mode": self.converter.mode}) + d[key] = self.converter(d[key]) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + for key in self.key_iterator(d): + tr = self.get_most_recent_transform(d, key) + original_mode = tr[TraceKeys.EXTRA_INFO]["mode"] + inverse_converter = ConvertBoxMode(src_mode=None, dst_mode=original_mode) + # Inverse is same as forward + d[key] = inverse_converter(d[key]) + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class AffineBoxToImageCoordinated(MapTransform, InvertibleTransform): + """ + Dictionary-based transform that converts box in world coordinate to image coordinate. + + Args: + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: The single key that represents the reference image to which ``box_keys`` are attached. + remove_empty: whether to remove the boxes that are actually empty + allow_missing_keys: don't raise exception if key is missing. + image_meta_key: explicitly indicate the key of the corresponding metadata dictionary. + for example, for data with key `image`, the metadata by default is in `image_meta_dict`. + the metadata is a dictionary object which contains: filename, affine, original_shape, etc. + it is a string, map to the `box_ref_image_key`. + if None, will try to construct meta_keys by `box_ref_image_key_{meta_key_postfix}`. + image_meta_key_postfix: if image_meta_keys=None, use `box_ref_image_key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. + For example, to handle key `image`, read/write affine matrices from the + metadata `image_meta_dict` dictionary's `affine` field. + affine_lps_to_ras: default ``False``. Yet if 1) the image is read by ITKReader, + and 2) the ITKReader has affine_lps_to_ras=True, and 3) the box is in world coordinate, + then set ``affine_lps_to_ras=True``. + """ + + def __init__( + self, + box_keys: KeysCollection, + box_ref_image_keys: str, + allow_missing_keys: bool = False, + image_meta_key: Union[str, None] = None, + image_meta_key_postfix: Union[str, None] = DEFAULT_POST_FIX, + affine_lps_to_ras=False, + ) -> None: + super().__init__(box_keys, allow_missing_keys) + box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys) + if len(box_ref_image_keys_tuple) > 1: + raise ValueError( + "Please provide a single key for box_ref_image_keys.\ + All boxes of box_keys are attached to box_ref_image_keys." + ) + self.image_meta_key = image_meta_key or f"{box_ref_image_keys}_{image_meta_key_postfix}" + self.converter_to_image_coordinate = AffineBox() + self.affine_lps_to_ras = affine_lps_to_ras + + def extract_affine(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + d = dict(data) + + meta_key = self.image_meta_key + # extract affine matrix from meta_data + if meta_key not in d: + raise ValueError(f"{meta_key} is not found. Please check whether it is the correct the image meta key.") + if "affine" not in d[meta_key]: + raise ValueError( + f"'affine' is not found in {meta_key}. \ + Please check whether it is the correct the image meta key." + ) + affine: NdarrayOrTensor = d[meta_key]["affine"] # type: ignore + if self.affine_lps_to_ras: # RAS affine + affine = orientation_ras_lps(affine) + + # when convert boxes from world coordinate to image coordinate, + # we apply inverse affine transform + affine_t, *_ = convert_data_type(affine, torch.Tensor) + # torch.inverse should not run in half precision + inv_affine_t = torch.inverse(affine_t.to(COMPUTE_DTYPE)) + return affine, inv_affine_t + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + + affine, inv_affine_t = self.extract_affine(data) + + for key in self.key_iterator(d): + self.push_transform(d, key, extra_info={"affine": affine}) + d[key] = self.converter_to_image_coordinate(d[key], affine=inv_affine_t) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + affine = transform["extra_info"]["affine"] + d[key] = AffineBox()(d[key], affine=affine) + self.pop_transform(d, key) + return d + + +class AffineBoxToWorldCoordinated(AffineBoxToImageCoordinated): + """ + Dictionary-based transform that converts box in image coordinate to world coordinate. + + Args: + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: The single key that represents the reference image to which ``box_keys`` are attached. + remove_empty: whether to remove the boxes that are actually empty + allow_missing_keys: don't raise exception if key is missing. + image_meta_key: explicitly indicate the key of the corresponding metadata dictionary. + for example, for data with key `image`, the metadata by default is in `image_meta_dict`. + the metadata is a dictionary object which contains: filename, affine, original_shape, etc. + it is a string, map to the `box_ref_image_key`. + if None, will try to construct meta_keys by `box_ref_image_key_{meta_key_postfix}`. + image_meta_key_postfix: if image_meta_keys=None, use `box_ref_image_key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. + For example, to handle key `image`, read/write affine matrices from the + metadata `image_meta_dict` dictionary's `affine` field. + affine_lps_to_ras: default ``False``. Yet if 1) the image is read by ITKReader, + and 2) the ITKReader has affine_lps_to_ras=True, and 3) the box is in world coordinate, + then set ``affine_lps_to_ras=True``. + """ + + def __init__( + self, + box_keys: KeysCollection, + box_ref_image_keys: str, + allow_missing_keys: bool = False, + image_meta_key: Union[str, None] = None, + image_meta_key_postfix: Union[str, None] = DEFAULT_POST_FIX, + affine_lps_to_ras=False, + ) -> None: + super().__init__( + box_keys, box_ref_image_keys, allow_missing_keys, image_meta_key, image_meta_key_postfix, affine_lps_to_ras + ) + self.converter_to_world_coordinate = AffineBox() + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + + affine, inv_affine_t = self.extract_affine(data) + + for key in self.key_iterator(d): + self.push_transform(d, key, extra_info={"affine": inv_affine_t}) + d[key] = self.converter_to_world_coordinate(d[key], affine=affine) + return d + + +class ZoomBoxd(MapTransform, InvertibleTransform): + """ + Dictionary-based transform that zooms input boxes and images with the given zoom scale. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + zoom: The zoom factor along the spatial axes. + If a float, zoom is the same for each spatial axis. + If a sequence, zoom should contain one value for each spatial axis. + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + The interpolation mode. Defaults to ``"area"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + It also can be a sequence of string, each element corresponds to a key in ``keys``. + padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + The mode to pad data after zooming. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + align_corners: This only has an effect when mode is + 'linear', 'bilinear', 'bicubic' or 'trilinear'. Default: None. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + It also can be a sequence of bool or None, each element corresponds to a key in ``keys``. + keep_size: Should keep original size (pad if needed), default is True. + allow_missing_keys: don't raise exception if key is missing. + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. + """ + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + zoom: Union[Sequence[float], float], + mode: SequenceStr = InterpolateMode.AREA, + padding_mode: SequenceStr = NumpyPadMode.EDGE, + align_corners: Union[Sequence[Optional[bool]], Optional[bool]] = None, + keep_size: bool = True, + allow_missing_keys: bool = False, + **kwargs, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + super().__init__(self.image_keys + self.box_keys, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.mode = ensure_tuple_rep(mode, len(self.image_keys)) + self.padding_mode = ensure_tuple_rep(padding_mode, len(self.image_keys)) + self.align_corners = ensure_tuple_rep(align_corners, len(self.image_keys)) + self.zoomer = Zoom(zoom=zoom, keep_size=keep_size, **kwargs) + self.keep_size = keep_size + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + + # zoom box + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + src_spatial_size = d[box_ref_image_key].shape[1:] + dst_spatial_size = [int(round(z * ss)) for z, ss in zip(self.zoomer.zoom, src_spatial_size)] # type: ignore + self.zoomer.zoom = [ds / float(ss) for ss, ds in zip(src_spatial_size, dst_spatial_size)] + self.push_transform( + d, + box_key, + extra_info={"zoom": self.zoomer.zoom, "src_spatial_size": src_spatial_size, "type": "box_key"}, + ) + d[box_key] = ZoomBox(zoom=self.zoomer.zoom, keep_size=self.keep_size)( + d[box_key], src_spatial_size=src_spatial_size + ) + + # zoom image, copied from monai.transforms.spatial.dictionary.Zoomd + for key, mode, padding_mode, align_corners in zip( + self.image_keys, self.mode, self.padding_mode, self.align_corners + ): + self.push_transform( + d, + key, + extra_info={ + "mode": mode.value if isinstance(mode, Enum) else mode, + "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, + "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, + "original_shape": d[key].shape[1:], + "type": "image_key", + }, + ) + d[key] = self.zoomer(d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners) + + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + # zoom image, copied from monai.transforms.spatial.dictionary.Zoomd + if key_type == "image_key": + # Create inverse transform + zoom = np.array(self.zoomer.zoom) + inverse_transform = Zoom(zoom=(1 / zoom).tolist(), keep_size=self.zoomer.keep_size) + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] + align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] + # Apply inverse + d[key] = inverse_transform( + d[key], + mode=mode, + padding_mode=padding_mode, + align_corners=None if align_corners == TraceKeys.NONE else align_corners, + ) + # Size might be out by 1 voxel so pad + orig_shape = transform[TraceKeys.EXTRA_INFO]["original_shape"] + d[key] = SpatialPad(orig_shape, mode="edge")(d[key]) + + # zoom boxes + if key_type == "box_key": + zoom = np.array(transform[TraceKeys.EXTRA_INFO]["zoom"]) + src_spatial_size = transform[TraceKeys.EXTRA_INFO]["src_spatial_size"] + box_inverse_transform = ZoomBox(zoom=(1 / zoom).tolist(), keep_size=self.zoomer.keep_size) + d[key] = box_inverse_transform(d[key], src_spatial_size=src_spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) + + return d + + +class RandZoomBoxd(RandomizableTransform, MapTransform, InvertibleTransform): + """ + Dictionary-based transform that randomly zooms input boxes and images with given probability within given zoom range. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + prob: Probability of zooming. + min_zoom: Min zoom factor. Can be float or sequence same size as image. + If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims + to keep the original spatial shape ratio. + If a sequence, min_zoom should contain one value for each spatial axis. + If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio. + max_zoom: Max zoom factor. Can be float or sequence same size as image. + If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims + to keep the original spatial shape ratio. + If a sequence, max_zoom should contain one value for each spatial axis. + If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio. + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + The interpolation mode. Defaults to ``"area"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + It also can be a sequence of string, each element corresponds to a key in ``keys``. + padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + The mode to pad data after zooming. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + align_corners: This only has an effect when mode is + 'linear', 'bilinear', 'bicubic' or 'trilinear'. Default: None. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + It also can be a sequence of bool or None, each element corresponds to a key in ``keys``. + keep_size: Should keep original size (pad if needed), default is True. + allow_missing_keys: don't raise exception if key is missing. + kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. + more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + """ + + backend = RandZoom.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + prob: float = 0.1, + min_zoom: Union[Sequence[float], float] = 0.9, + max_zoom: Union[Sequence[float], float] = 1.1, + mode: SequenceStr = InterpolateMode.AREA, + padding_mode: SequenceStr = NumpyPadMode.EDGE, + align_corners: Union[Sequence[Optional[bool]], Optional[bool]] = None, + keep_size: bool = True, + allow_missing_keys: bool = False, + **kwargs, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + MapTransform.__init__(self, self.image_keys + self.box_keys, allow_missing_keys) + RandomizableTransform.__init__(self, prob) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.rand_zoom = RandZoom(prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, keep_size=keep_size, **kwargs) + self.mode = ensure_tuple_rep(mode, len(self.image_keys)) + self.padding_mode = ensure_tuple_rep(padding_mode, len(self.image_keys)) + self.align_corners = ensure_tuple_rep(align_corners, len(self.image_keys)) + self.keep_size = keep_size + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandZoomBoxd": + super().set_random_state(seed, state) + self.rand_zoom.set_random_state(seed, state) + return self + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + first_key: Union[Hashable, List] = self.first_key(d) + if first_key == []: + return d + + self.randomize(None) + + # all the keys share the same random zoom factor + self.rand_zoom.randomize(d[first_key]) # type: ignore + + # zoom box + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + if self._do_transform: + src_spatial_size = d[box_ref_image_key].shape[1:] + dst_spatial_size = [int(round(z * ss)) for z, ss in zip(self.rand_zoom._zoom, src_spatial_size)] + self.rand_zoom._zoom = [ds / float(ss) for ss, ds in zip(src_spatial_size, dst_spatial_size)] + + self.push_transform( + d, + box_key, + extra_info={"zoom": self.rand_zoom._zoom, "src_spatial_size": src_spatial_size, "type": "box_key"}, + ) + d[box_key] = ZoomBox(zoom=self.rand_zoom._zoom, keep_size=self.keep_size)( + d[box_key], src_spatial_size=src_spatial_size + ) + + # zoom image, copied from monai.transforms.spatial.dictionary.RandZoomd + for key, mode, padding_mode, align_corners in zip( + self.image_keys, self.mode, self.padding_mode, self.align_corners + ): + if self._do_transform: + self.push_transform( + d, + key, + extra_info={ + "zoom": self.rand_zoom._zoom, + "mode": mode.value if isinstance(mode, Enum) else mode, + "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, + "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, + "original_shape": d[key].shape[1:], + "type": "image_key", + }, + ) + d[key] = self.rand_zoom( + d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners, randomize=False + ) + + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + # Check if random transform was actually performed (based on `prob`) + if transform[TraceKeys.DO_TRANSFORM]: + # zoom image, copied from monai.transforms.spatial.dictionary.Zoomd + if key_type == "image_key": + zoom = np.array(transform[TraceKeys.EXTRA_INFO]["zoom"]) + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] + align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] + inverse_transform = Zoom(zoom=(1.0 / zoom).tolist(), keep_size=self.rand_zoom.keep_size) + d[key] = inverse_transform( + d[key], + mode=mode, + padding_mode=padding_mode, + align_corners=None if align_corners == TraceKeys.NONE else align_corners, + ) + # Size might be out by 1 voxel so pad + orig_shape = transform[TraceKeys.EXTRA_INFO]["original_shape"] + d[key] = SpatialPad(orig_shape, mode="edge")(d[key]) + + # zoom boxes + if key_type == "box_key": + # Create inverse transform + zoom = np.array(transform[TraceKeys.EXTRA_INFO]["zoom"]) + src_spatial_size = transform[TraceKeys.EXTRA_INFO]["src_spatial_size"] + box_inverse_transform = ZoomBox(zoom=(1.0 / zoom).tolist(), keep_size=self.rand_zoom.keep_size) + d[key] = box_inverse_transform(d[key], src_spatial_size=src_spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class FlipBoxd(MapTransform, InvertibleTransform): + """ + Dictionary-based transform that flip boxes and images. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + spatial_axis: Spatial axes along which to flip over. Default is None. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = Flip.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + spatial_axis: Optional[Union[Sequence[int], int]] = None, + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + super().__init__(self.image_keys + self.box_keys, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.flipper = Flip(spatial_axis=spatial_axis) + self.box_flipper = FlipBox(spatial_axis=self.flipper.spatial_axis) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + + for key in self.image_keys: + d[key] = self.flipper(d[key]) + self.push_transform(d, key, extra_info={"type": "image_key"}) + + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + spatial_size = d[box_ref_image_key].shape[1:] + d[box_key] = self.box_flipper(d[box_key], spatial_size) + self.push_transform(d, box_key, extra_info={"spatial_size": spatial_size, "type": "box_key"}) + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + + # flip image, copied from monai.transforms.spatial.dictionary.Flipd + if key_type == "image_key": + d[key] = self.flipper(d[key]) + + # flip boxes + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + d[key] = self.box_flipper(d[key], spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class RandFlipBoxd(RandomizableTransform, MapTransform, InvertibleTransform): + """ + Dictionary-based transform that randomly flip boxes and images with the given probabilities. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + prob: Probability of flipping. + spatial_axis: Spatial axes along which to flip over. Default is None. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = RandFlip.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + prob: float = 0.1, + spatial_axis: Optional[Union[Sequence[int], int]] = None, + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + MapTransform.__init__(self, self.image_keys + self.box_keys, allow_missing_keys) + RandomizableTransform.__init__(self, prob) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.flipper = RandFlip(prob=1.0, spatial_axis=spatial_axis) + self.box_flipper = FlipBox(spatial_axis=spatial_axis) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandFlipBoxd": + super().set_random_state(seed, state) + self.flipper.set_random_state(seed, state) + return self + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + self.randomize(None) + + for key in self.image_keys: + if self._do_transform: + d[key] = self.flipper(d[key], randomize=False) + self.push_transform(d, key, extra_info={"type": "image_key"}) + + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + spatial_size = d[box_ref_image_key].shape[1:] + if self._do_transform: + d[box_key] = self.box_flipper(d[box_key], spatial_size) + self.push_transform(d, box_key, extra_info={"spatial_size": spatial_size, "type": "box_key"}) + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + # Check if random transform was actually performed (based on `prob`) + if transform[TraceKeys.DO_TRANSFORM]: + # flip image, copied from monai.transforms.spatial.dictionary.RandFlipd + if key_type == "image_key": + d[key] = self.flipper(d[key], randomize=False) + + # flip boxes + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + d[key] = self.box_flipper(d[key], spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class ClipBoxToImaged(MapTransform): + """ + Dictionary-based wrapper of :py:class:`monai.apps.detection.transforms.array.ClipBoxToImage`. + + Clip the bounding boxes and the associated labels/scores to makes sure they are within the image. + There might be multiple keys of labels/scores associated with one key of boxes. + + Args: + box_keys: The single key to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + label_keys: Keys that represent the labels corresponding to the ``box_keys``. Multiple keys are allowed. + box_ref_image_keys: The single key that represents the reference image + to which ``box_keys`` and ``label_keys`` are attached. + remove_empty: whether to remove the boxes that are actually empty + allow_missing_keys: don't raise exception if key is missing. + + Example: + .. code-block:: python + + ClipBoxToImaged( + box_keys="boxes", box_ref_image_keys="image", label_keys=["labels", "scores"], remove_empty=True + ) + """ + + def __init__( + self, + box_keys: KeysCollection, + label_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + remove_empty: bool = True, + allow_missing_keys: bool = False, + ) -> None: + box_keys_tuple = ensure_tuple(box_keys) + if len(box_keys_tuple) != 1: + raise ValueError( + "Please provide a single key for box_keys.\ + All label_keys are attached to this box_keys." + ) + box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys) + if len(box_ref_image_keys_tuple) != 1: + raise ValueError( + "Please provide a single key for box_ref_image_keys.\ + All box_keys and label_keys are attached to this box_ref_image_keys." + ) + self.label_keys = ensure_tuple(label_keys) + super().__init__(box_keys_tuple, allow_missing_keys) + + self.box_keys = box_keys_tuple[0] + self.box_ref_image_keys = box_ref_image_keys_tuple[0] + self.clipper = ClipBoxToImage(remove_empty=remove_empty) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + spatial_size = d[self.box_ref_image_keys].shape[1:] + labels = [d[label_key] for label_key in self.label_keys] # could be multiple arrays + d[self.box_keys], clipped_labels = self.clipper(d[self.box_keys], labels, spatial_size) + + for label_key, clipped_labels_i in zip(self.label_keys, clipped_labels): + d[label_key] = clipped_labels_i + return d + + +class BoxToMaskd(MapTransform): + """ + Dictionary-based wrapper of :py:class:`monai.apps.detection.transforms.array.BoxToMask`. + Pairs with :py:class:`monai.apps.detection.transforms.dictionary.MaskToBoxd` . + Please make sure the same ``min_fg_label`` is used when using the two transforms in pairs. + The output ``d[box_mask_key]`` will have background intensity 0, since the following operations + may pad 0 on the border. + + This is the general solution for transforms that need to be applied on images and boxes simultaneously. + It is performed with the following steps. + + 1) use ``BoxToMaskd`` to covert boxes and labels to box_masks; + 2) do transforms, e.g., rotation or cropping, on images and box_masks together; + 3) use ``MaskToBoxd`` to convert box_masks back to boxes and labels. + + Args: + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_mask_keys: Keys to store output box mask results for transformation. Same length with ``box_keys``. + label_keys: Keys that represent the labels corresponding to the ``box_keys``. Same length with ``box_keys``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + min_fg_label: min foreground box label. + ellipse_mask: bool. + + - If True, it assumes the object shape is close to ellipse or ellipsoid. + - If False, it assumes the object shape is close to rectangle or cube and well occupies the bounding box. + - If the users are going to apply random rotation as data augmentation, we suggest setting ellipse_mask=True + See also Kalra et al. "Towards Rotation Invariance in Object Detection", ICCV 2021. + allow_missing_keys: don't raise exception if key is missing. + + Example: + .. code-block:: python + + # This code snippet creates transforms (random rotation and cropping) on boxes, labels, and image together. + import numpy as np + from monai.transforms import Compose, RandRotated, RandSpatialCropd, DeleteItemsd + transforms = Compose( + [ + BoxToMaskd( + box_keys="boxes", label_keys="labels", + box_mask_keys="box_mask", box_ref_image_keys="image", + min_fg_label=0, ellipse_mask=True + ), + RandRotated(keys=["image","box_mask"],mode=["nearest","nearest"], + prob=0.2,range_x=np.pi/6,range_y=np.pi/6,range_z=np.pi/6, + keep_size=True,padding_mode="zeros" + ), + RandSpatialCropd(keys=["image","box_mask"],roi_size=128, random_size=False), + MaskToBoxd( + box_mask_keys="box_mask", box_keys="boxes", + label_keys="labels", min_fg_label=0 + ) + DeleteItemsd(keys=["box_mask"]), + ] + ) + + """ + + def __init__( + self, + box_keys: KeysCollection, + box_mask_keys: KeysCollection, + label_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + min_fg_label: int, + ellipse_mask: bool = False, + allow_missing_keys: bool = False, + ) -> None: + super().__init__(box_keys, allow_missing_keys) + self.box_keys = ensure_tuple(box_keys) + self.label_keys = ensure_tuple(label_keys) + self.box_mask_keys = ensure_tuple(box_mask_keys) + if not len(self.label_keys) == len(self.box_keys) == len(self.box_mask_keys): + raise ValueError("Please make sure len(label_keys)==len(box_keys)==len(box_mask_keys)!") + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + self.bg_label = min_fg_label - 1 # make sure background label is always smaller than fg labels. + self.converter = BoxToMask(bg_label=self.bg_label, ellipse_mask=ellipse_mask) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + + for box_key, label_key, box_mask_key, box_ref_image_key in zip( + self.box_keys, self.label_keys, self.box_mask_keys, self.box_ref_image_keys + ): + spatial_size = d[box_ref_image_key].shape[1:] + d[box_mask_key] = self.converter(d[box_key], d[label_key], spatial_size) + # make box mask background intensity to be 0, since the following operations may pad 0 on the border. + d[box_mask_key] -= self.bg_label + return d + + +class MaskToBoxd(MapTransform): + """ + Dictionary-based wrapper of :py:class:`monai.apps.detection.transforms.array.MaskToBox`. + Pairs with :py:class:`monai.apps.detection.transforms.dictionary.BoxToMaskd` . + Please make sure the same ``min_fg_label`` is used when using the two transforms in pairs. + + This is the general solution for transforms that need to be applied on images and boxes simultaneously. + It is performed with the following steps. + + 1) use ``BoxToMaskd`` to covert boxes and labels to box_masks; + 2) do transforms, e.g., rotation or cropping, on images and box_masks together; + 3) use ``MaskToBoxd`` to convert box_masks back to boxes and labels. + + Args: + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_mask_keys: Keys to store output box mask results for transformation. Same length with ``box_keys``. + label_keys: Keys that represent the labels corresponding to the ``box_keys``. Same length with ``box_keys``. + min_fg_label: min foreground box label. + box_dtype: output dtype for box_keys + label_dtype: output dtype for label_keys + allow_missing_keys: don't raise exception if key is missing. + + Example: + .. code-block:: python + + # This code snippet creates transforms (random rotation and cropping) on boxes, labels, and images together. + import numpy as np + from monai.transforms import Compose, RandRotated, RandSpatialCropd, DeleteItemsd + transforms = Compose( + [ + BoxToMaskd( + box_keys="boxes", label_keys="labels", + box_mask_keys="box_mask", box_ref_image_keys="image", + min_fg_label=0, ellipse_mask=True + ), + RandRotated(keys=["image","box_mask"],mode=["nearest","nearest"], + prob=0.2,range_x=np.pi/6,range_y=np.pi/6,range_z=np.pi/6, + keep_size=True,padding_mode="zeros" + ), + RandSpatialCropd(keys=["image","box_mask"],roi_size=128, random_size=False), + MaskToBoxd( + box_mask_keys="box_mask", box_keys="boxes", + label_keys="labels", min_fg_label=0 + ) + DeleteItemsd(keys=["box_mask"]), + ] + ) + """ + + def __init__( + self, + box_keys: KeysCollection, + box_mask_keys: KeysCollection, + label_keys: KeysCollection, + min_fg_label: int, + box_dtype=torch.float32, + label_dtype=torch.long, + allow_missing_keys: bool = False, + ) -> None: + super().__init__(box_keys, allow_missing_keys) + self.box_keys = ensure_tuple(box_keys) + self.label_keys = ensure_tuple(label_keys) + self.box_mask_keys = ensure_tuple(box_mask_keys) + if not len(self.label_keys) == len(self.box_keys) == len(self.box_mask_keys): + raise ValueError("Please make sure len(label_keys)==len(box_keys)==len(box_mask_keys)!") + self.bg_label = min_fg_label - 1 # make sure background label is always smaller than fg labels. + self.converter = MaskToBox(bg_label=self.bg_label, box_dtype=box_dtype, label_dtype=label_dtype) + self.box_dtype = box_dtype + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + + for box_key, label_key, box_mask_key in zip(self.box_keys, self.label_keys, self.box_mask_keys): + d[box_mask_key] += self.bg_label # pairs with the operation in BoxToMaskd + d[box_key], d[label_key] = self.converter(d[box_mask_key]) + return d + + +class RandCropBoxByPosNegLabeld(Randomizable, MapTransform): + """ + Crop random fixed sized regions that contains foreground boxes. + Suppose all the expected fields specified by `image_keys` have same shape, + and add `patch_index` to the corresponding meta data. + And will return a list of dictionaries for all the cropped images. + If a dimension of the expected spatial size is bigger than the input image size, + will not crop that dimension. So the cropped result may be smaller than the expected size, + and the cropped results of several images may not have exactly the same shape. + + Args: + image_keys: Keys to pick image data for transformation. They need to have the same spatial size. + box_keys: The single key to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + label_keys: Keys that represent the labels corresponding to the ``box_keys``. Multiple keys are allowed. + spatial_size: the spatial size of the crop region e.g. [224, 224, 128]. + if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if its components have non-positive values, the corresponding size of `data[label_key]` will be used. + for example: if the spatial size of input data is [40, 40, 40] and `spatial_size=[32, 64, -1]`, + the spatial size of output data will be [32, 40, 40]. + pos: used with `neg` together to calculate the ratio ``pos / (pos + neg)`` for the probability + to pick a foreground voxel as a center rather than a background voxel. + neg: used with `pos` together to calculate the ratio ``pos / (pos + neg)`` for the probability + to pick a foreground voxel as a center rather than a background voxel. + num_samples: number of samples (crop regions) to take in each list. + whole_box: Bool, default True, whether we prefer to contain at least one whole box in the cropped foreground patch. + Even if True, it is still possible to get partial box if there are multiple boxes in the image. + thresh_image_key: if thresh_image_key is not None, use ``label == 0 & thresh_image > image_threshold`` to select + the negative sample(background) center. so the crop center will only exist on valid image area. + image_threshold: if enabled thresh_image_key, use ``thresh_image > image_threshold`` to determine + the valid image content area. + fg_indices_key: if provided pre-computed foreground indices of `label`, will ignore above `image_key` and + `image_threshold`, and randomly select crop centers based on them, need to provide `fg_indices_key` + and `bg_indices_key` together, expect to be 1 dim array of spatial indices after flattening. + a typical usage is to call `FgBgToIndicesd` transform first and cache the results. + bg_indices_key: if provided pre-computed background indices of `label`, will ignore above `image_key` and + `image_threshold`, and randomly select crop centers based on them, need to provide `fg_indices_key` + and `bg_indices_key` together, expect to be 1 dim array of spatial indices after flattening. + a typical usage is to call `FgBgToIndicesd` transform first and cache the results. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. + used to add `patch_index` to the meta dict. + for example, for data with key `image`, the metadata by default is in `image_meta_dict`. + the metadata is a dictionary object which contains: filename, original_shape, etc. + it can be a sequence of string, map to the `keys`. + if None, will try to construct meta_keys by `key_{meta_key_postfix}`. + meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. + used to add `patch_index` to the meta dict. + allow_smaller: if `False`, an exception will be raised if the image is smaller than + the requested ROI in any dimension. If `True`, any smaller dimensions will be set to + match the cropped size (i.e., no cropping in that dimension). + allow_missing_keys: don't raise exception if key is missing. + """ + + def __init__( + self, + image_keys: KeysCollection, + box_keys: str, + label_keys: KeysCollection, + spatial_size: Union[Sequence[int], int], + pos: float = 1.0, + neg: float = 1.0, + num_samples: int = 1, + whole_box: bool = True, + thresh_image_key: Optional[str] = None, + image_threshold: float = 0.0, + fg_indices_key: Optional[str] = None, + bg_indices_key: Optional[str] = None, + meta_keys: Optional[KeysCollection] = None, + meta_key_postfix: str = DEFAULT_POST_FIX, + allow_smaller: bool = False, + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + if len(self.image_keys) < 1: + raise ValueError("At least one image_keys should be provided.") + + MapTransform.__init__(self, self.image_keys, allow_missing_keys) + + box_keys_tuple = ensure_tuple(box_keys) + if len(box_keys_tuple) != 1: + raise ValueError( + "Please provide a single key for box_keys.\ + All label_keys are attached to this box_keys." + ) + self.box_keys = box_keys_tuple[0] + self.label_keys = ensure_tuple(label_keys) + + self.spatial_size_: Union[Tuple[int, ...], Sequence[int], int] = spatial_size + + if pos < 0 or neg < 0: + raise ValueError(f"pos and neg must be nonnegative, got pos={pos} neg={neg}.") + if pos + neg == 0: + raise ValueError("Incompatible values: pos=0 and neg=0.") + self.pos_ratio = pos / (pos + neg) + if num_samples < 1: + raise ValueError(f"num_samples needs to be positive int, got num_samples={num_samples}.") + self.num_samples = num_samples + self.whole_box = whole_box + + self.thresh_image_key = thresh_image_key + self.image_threshold = image_threshold + self.fg_indices_key = fg_indices_key + self.bg_indices_key = bg_indices_key + + self.meta_keys = ensure_tuple_rep(None, len(self.image_keys)) if meta_keys is None else ensure_tuple(meta_keys) + if len(self.image_keys) != len(self.meta_keys): + raise ValueError("meta_keys should have the same length as keys.") + self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.image_keys)) + self.centers: Optional[List[List[int]]] = None + self.allow_smaller = allow_smaller + + def generate_fg_center_boxes_np(self, boxes: NdarrayOrTensor, image_size: Sequence[int]) -> np.ndarray: + # We don't require crop center to be whthin the boxes. + # As along as the cropped patch contains a box, it is considered as a foreground patch. + # Positions within extended_boxes are crop centers for foreground patches + spatial_dims = len(image_size) + boxes_np, *_ = convert_data_type(boxes, np.ndarray) + + extended_boxes = np.zeros_like(boxes_np, dtype=int) + boxes_start = np.ceil(boxes_np[:, :spatial_dims]).astype(int) + boxes_stop = np.floor(boxes_np[:, spatial_dims:]).astype(int) + for axis in range(spatial_dims): + if not self.whole_box: + extended_boxes[:, axis] = boxes_start[:, axis] - self.spatial_size[axis] // 2 + 1 + extended_boxes[:, axis + spatial_dims] = boxes_stop[:, axis] + self.spatial_size[axis] // 2 - 1 + else: + # extended box start + extended_boxes[:, axis] = boxes_stop[:, axis] - self.spatial_size[axis] // 2 - 1 + extended_boxes[:, axis] = np.minimum(extended_boxes[:, axis], boxes_start[:, axis]) + # extended box stop + extended_boxes[:, axis + spatial_dims] = extended_boxes[:, axis] + self.spatial_size[axis] // 2 + extended_boxes[:, axis + spatial_dims] = np.maximum( + extended_boxes[:, axis + spatial_dims], boxes_stop[:, axis] + ) + extended_boxes, _ = clip_boxes_to_image(extended_boxes, image_size, remove_empty=True) # type: ignore + return extended_boxes + + def randomize( # type: ignore + self, + boxes: NdarrayOrTensor, + image_size: Sequence[int], + fg_indices: Optional[NdarrayOrTensor] = None, + bg_indices: Optional[NdarrayOrTensor] = None, + thresh_image: Optional[NdarrayOrTensor] = None, + ) -> None: + if fg_indices is None or bg_indices is None: + # We don't require crop center to be whthin the boxes. + # As along as the cropped patch contains a box, it is considered as a foreground patch. + # Positions within extended_boxes are crop centers for foreground patches + extended_boxes_np = self.generate_fg_center_boxes_np(boxes, image_size) + mask_img = convert_box_to_mask( + extended_boxes_np, np.ones(extended_boxes_np.shape[0]), image_size, bg_label=0, ellipse_mask=False + ) + mask_img = np.amax(mask_img, axis=0, keepdims=True)[0:1, ...] + fg_indices_, bg_indices_ = map_binary_to_indices(mask_img, thresh_image, self.image_threshold) + else: + fg_indices_ = fg_indices + bg_indices_ = bg_indices + + self.centers = generate_pos_neg_label_crop_centers( + self.spatial_size, + self.num_samples, + self.pos_ratio, + image_size, + fg_indices_, + bg_indices_, + self.R, + self.allow_smaller, + ) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> List[Dict[Hashable, torch.Tensor]]: + d = dict(data) + spatial_dims = len(d[self.image_keys[0]].shape) - 1 + image_size = d[self.image_keys[0]].shape[1:] + self.spatial_size = ensure_tuple_rep(self.spatial_size_, spatial_dims) + + # randomly sample crop centers + boxes = d[self.box_keys] + labels = [d[label_key] for label_key in self.label_keys] # could be multiple arrays + fg_indices = d.pop(self.fg_indices_key, None) if self.fg_indices_key is not None else None + bg_indices = d.pop(self.bg_indices_key, None) if self.bg_indices_key is not None else None + thresh_image = d[self.thresh_image_key] if self.thresh_image_key else None + self.randomize(boxes, image_size, fg_indices, bg_indices, thresh_image) + + if self.centers is None: + raise ValueError("no available ROI centers to crop.") + + # initialize returned list with shallow copy to preserve key ordering + results: List[Dict[Hashable, torch.Tensor]] = [dict(d) for _ in range(self.num_samples)] + + # crop images and boxes for each center. + for i, center in enumerate(self.centers): + results[i] = deepcopy(d) + # compute crop start and end, always crop, no padding + cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) + crop_start = [max(s.start, 0) for s in cropper.slices] + crop_end = [min(s.stop, image_size_a) for s, image_size_a in zip(cropper.slices, image_size)] + crop_slices = [slice(int(s), int(e)) for s, e in zip(crop_start, crop_end)] + + # crop images + cropper = SpatialCrop(roi_slices=crop_slices) + for image_key in self.image_keys: + results[i][image_key] = cropper(d[image_key]) + + # crop boxes and labels + boxcropper = SpatialCropBox(roi_slices=crop_slices) + results[i][self.box_keys], cropped_labels = boxcropper(boxes, labels) + for label_key, cropped_labels_i in zip(self.label_keys, cropped_labels): + results[i][label_key] = cropped_labels_i + + # add `patch_index` to the meta data + for key, meta_key, meta_key_postfix in zip(self.image_keys, self.meta_keys, self.meta_key_postfix): + meta_key = meta_key or f"{key}_{meta_key_postfix}" + if meta_key not in results[i]: + results[i][meta_key] = {} # type: ignore + results[i][meta_key][Key.PATCH_INDEX] = i # type: ignore + + return results + + +class RotateBox90d(MapTransform, InvertibleTransform): + """ + Input boxes and images are rotated by 90 degrees + in the plane specified by ``spatial_axes`` for ``k`` times + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + k: number of times to rotate by 90 degrees. + spatial_axes: 2 int numbers, defines the plane to rotate with 2 spatial axes. + Default (0, 1), this is the first two axis in spatial dimensions. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = RotateBox90.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + k: int = 1, + spatial_axes: Tuple[int, int] = (0, 1), + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + super().__init__(self.image_keys + self.box_keys, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + self.img_rotator = Rotate90(k, spatial_axes) + self.box_rotator = RotateBox90(k, spatial_axes) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, torch.Tensor]: + d = dict(data) + for key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + spatial_size = list(d[box_ref_image_key].shape[1:]) + d[key] = self.box_rotator(d[key], spatial_size) + if self.img_rotator.k % 2 == 1: + # if k = 1 or 3, spatial_size will be transposed + spatial_size[self.img_rotator.spatial_axes[0]], spatial_size[self.img_rotator.spatial_axes[1]] = ( + spatial_size[self.img_rotator.spatial_axes[1]], + spatial_size[self.img_rotator.spatial_axes[0]], + ) + self.push_transform(d, key, extra_info={"spatial_size": spatial_size, "type": "box_key"}) + + for key in self.image_keys: + d[key] = self.img_rotator(d[key]) + self.push_transform(d, key, extra_info={"type": "image_key"}) + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + num_times_to_rotate = 4 - self.img_rotator.k + + if key_type == "image_key": + inverse_transform = Rotate90(num_times_to_rotate, self.img_rotator.spatial_axes) + d[key] = inverse_transform(d[key]) + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + inverse_transform = RotateBox90(num_times_to_rotate, self.box_rotator.spatial_axes) + d[key] = inverse_transform(d[key], spatial_size) + self.pop_transform(d, key) + return d + + +class RandRotateBox90d(RandRotate90d): + """ + With probability `prob`, input boxes and images are rotated by 90 degrees + in the plane specified by `spatial_axes`. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represent the reference images to which ``box_keys`` are attached. + prob: probability of rotating. + (Default 0.1, with 10% probability it returns a rotated array.) + max_k: number of rotations will be sampled from `np.random.randint(max_k) + 1`. + (Default 3) + spatial_axes: 2 int numbers, defines the plane to rotate with 2 spatial axes. + Default: (0, 1), this is the first two axis in spatial dimensions. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = RotateBox90.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + prob: float = 0.1, + max_k: int = 3, + spatial_axes: Tuple[int, int] = (0, 1), + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + super().__init__(self.image_keys + self.box_keys, prob, max_k, spatial_axes, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, torch.Tensor]: + self.randomize() + d = dict(data) + + if self._rand_k % 4 == 0: + return d + + # FIXME: here we didn't use array version `RandRotate90` transform as others, because we need + # to be compatible with the random status of some previous integration tests + box_rotator = RotateBox90(self._rand_k, self.spatial_axes) + img_rotator = Rotate90(self._rand_k, self.spatial_axes) + + for key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + if self._do_transform: + spatial_size = list(d[box_ref_image_key].shape[1:]) + d[key] = box_rotator(d[key], spatial_size) + if self._rand_k % 2 == 1: + # if k = 1 or 3, spatial_size will be transposed + spatial_size[self.spatial_axes[0]], spatial_size[self.spatial_axes[1]] = ( + spatial_size[self.spatial_axes[1]], + spatial_size[self.spatial_axes[0]], + ) + self.push_transform( + d, key, extra_info={"rand_k": self._rand_k, "spatial_size": spatial_size, "type": "box_key"} + ) + + for key in self.image_keys: + if self._do_transform: + d[key] = img_rotator(d[key]) + self.push_transform(d, key, extra_info={"rand_k": self._rand_k, "type": "image_key"}) + return d + + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + if self._rand_k % 4 == 0: + return d + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + # Check if random transform was actually performed (based on `prob`) + if transform[TraceKeys.DO_TRANSFORM]: + num_times_rotated = transform[TraceKeys.EXTRA_INFO]["rand_k"] + num_times_to_rotate = 4 - num_times_rotated + # flip image, copied from monai.transforms.spatial.dictionary.RandFlipd + if key_type == "image_key": + inverse_transform = Rotate90(num_times_to_rotate, self.spatial_axes) + d[key] = inverse_transform(d[key]) + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + inverse_transform = RotateBox90(num_times_to_rotate, self.spatial_axes) + d[key] = inverse_transform(d[key], spatial_size) + self.pop_transform(d, key) + return d + + +ConvertBoxModeD = ConvertBoxModeDict = ConvertBoxModed +ConvertBoxToStandardModeD = ConvertBoxToStandardModeDict = ConvertBoxToStandardModed +ZoomBoxD = ZoomBoxDict = ZoomBoxd +RandZoomBoxD = RandZoomBoxDict = RandZoomBoxd +AffineBoxToImageCoordinateD = AffineBoxToImageCoordinateDict = AffineBoxToImageCoordinated +FlipBoxD = FlipBoxDict = FlipBoxd +RandFlipBoxD = RandFlipBoxDict = RandFlipBoxd +ClipBoxToImageD = ClipBoxToImageDict = ClipBoxToImaged +BoxToMaskD = BoxToMaskDict = BoxToMaskd +MaskToBoxD = MaskToBoxDict = MaskToBoxd +RandCropBoxByPosNegLabelD = RandCropBoxByPosNegLabelDict = RandCropBoxByPosNegLabeld +RotateBox90D = RotateBox90Dict = RotateBox90d +RandRotateBox90D = RandRotateBox90Dict = RandRotateBox90d diff --git a/monai/apps/detection/utils/ATSS_matcher.py b/monai/apps/detection/utils/ATSS_matcher.py new file mode 100644 index 0000000000..f0170422bb --- /dev/null +++ b/monai/apps/detection/utils/ATSS_matcher.py @@ -0,0 +1,288 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/matcher.py +# which has the following license... +# https://github.com/MIC-DKFZ/nnDetection/blob/main/LICENSE +# +# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/_utils.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE +# +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +The functions in this script are adapted from nnDetection, +https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/matcher.py +which is adapted from torchvision. + +These are the changes compared with nndetection: +1) comments and docstrings; +2) reformat; +3) add a debug option to ATSSMatcher to help the users to tune parameters; +4) add a corner case return in ATSSMatcher.compute_matches; +5) add support for float16 cpu +""" + +import logging +from abc import ABC +from typing import Callable, Sequence, Tuple, TypeVar + +import torch +from torch import Tensor + +from monai.data.box_utils import COMPUTE_DTYPE, box_iou, boxes_center_distance, centers_in_boxes +from monai.utils.type_conversion import convert_to_tensor + +# -INF should be smaller than the lower bound of similarity_fn output. +INF = float("inf") + + +class Matcher(ABC): + """ + Base class of Matcher, which matches boxes and anchors to each other + + Args: + similarity_fn: function for similarity computation between + boxes and anchors + """ + + BELOW_LOW_THRESHOLD: int = -1 + BETWEEN_THRESHOLDS: int = -2 + + def __init__(self, similarity_fn: Callable[[Tensor, Tensor], Tensor] = box_iou): # type: ignore + self.similarity_fn = similarity_fn + + def __call__( + self, boxes: torch.Tensor, anchors: torch.Tensor, num_anchors_per_level: Sequence[int], num_anchors_per_loc: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute matches for a single image + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. + num_anchors_per_level: number of anchors per feature pyramid level + num_anchors_per_loc: number of anchors per position + + Returns: + - matrix which contains the similarity from each boxes to each anchor [N, M] + - vector which contains the matched box index for all + anchors (if background `BELOW_LOW_THRESHOLD` is used + and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] + + Note: + ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, + also represented as "xyxy" ([xmin, ymin, xmax, ymax]) for 2D + and "xyzxyz" ([xmin, ymin, zmin, xmax, ymax, zmax]) for 3D. + """ + if boxes.numel() == 0: + # no ground truth + num_anchors = anchors.shape[0] + match_quality_matrix = torch.tensor([]).to(anchors) + matches = torch.empty(num_anchors, dtype=torch.int64).fill_(self.BELOW_LOW_THRESHOLD) + return match_quality_matrix, matches + # at least one ground truth + return self.compute_matches( + boxes=boxes, + anchors=anchors, + num_anchors_per_level=num_anchors_per_level, + num_anchors_per_loc=num_anchors_per_loc, + ) + + def compute_matches( + self, boxes: torch.Tensor, anchors: torch.Tensor, num_anchors_per_level: Sequence[int], num_anchors_per_loc: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute matches + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. + num_anchors_per_level: number of anchors per feature pyramid level + num_anchors_per_loc: number of anchors per position + + Returns: + - matrix which contains the similarity from each boxes to each anchor [N, M] + - vector which contains the matched box index for all + anchors (if background `BELOW_LOW_THRESHOLD` is used + and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] + """ + raise NotImplementedError + + +class ATSSMatcher(Matcher): + def __init__( + self, + num_candidates: int = 4, + similarity_fn: Callable[[Tensor, Tensor], Tensor] = box_iou, # type: ignore + center_in_gt: bool = True, + debug: bool = False, + ): + """ + Compute matching based on ATSS https://arxiv.org/abs/1912.02424 + `Bridging the Gap Between Anchor-based and Anchor-free Detection + via Adaptive Training Sample Selection` + + Args: + num_candidates: number of positions to select candidates from. + Smaller value will result in a higher matcher threshold and less matched candidates. + similarity_fn: function for similarity computation between boxes and anchors + center_in_gt: If False (default), matched anchor center points do not need + to lie withing the ground truth box. Recommend False for small objects. + If True, will result in a strict matcher and less matched candidates. + debug: if True, will print the matcher threshold in order to + tune ``num_candidates`` and ``center_in_gt``. + """ + super().__init__(similarity_fn=similarity_fn) + self.num_candidates = num_candidates + self.min_dist = 0.01 + self.center_in_gt = center_in_gt + self.debug = debug + logging.info( + f"Running ATSS Matching with num_candidates={self.num_candidates} and center_in_gt {self.center_in_gt}." + ) + + def compute_matches( + self, boxes: torch.Tensor, anchors: torch.Tensor, num_anchors_per_level: Sequence[int], num_anchors_per_loc: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute matches according to ATTS for a single image + Adapted from + (https://github.com/sfzhang15/ATSS/blob/79dfb28bd1/atss_core/modeling/rpn/atss/loss.py#L180-L184) + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. + num_anchors_per_level: number of anchors per feature pyramid level + num_anchors_per_loc: number of anchors per position + + Returns: + - matrix which contains the similarity from each boxes to each anchor [N, M] + - vector which contains the matched box index for all + anchors (if background `BELOW_LOW_THRESHOLD` is used + and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] + + Note: + ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, + also represented as "xyxy" ([xmin, ymin, xmax, ymax]) for 2D + and "xyzxyz" ([xmin, ymin, zmin, xmax, ymax, zmax]) for 3D. + """ + num_gt = boxes.shape[0] + num_anchors = anchors.shape[0] + + distances_, _, anchors_center = boxes_center_distance(boxes, anchors) # num_boxes x anchors + distances = convert_to_tensor(distances_) + + # select candidates based on center distance + candidate_idx_list = [] + start_idx = 0 + for _, apl in enumerate(num_anchors_per_level): + end_idx = start_idx + apl * num_anchors_per_loc + + # topk: total number of candidates per position + topk = min(self.num_candidates * num_anchors_per_loc, apl) + # torch.topk() does not support float16 cpu, need conversion to float32 or float64 + _, idx = distances[:, start_idx:end_idx].to(COMPUTE_DTYPE).topk(topk, dim=1, largest=False) + # idx: shape [num_boxes x topk] + candidate_idx_list.append(idx + start_idx) + + start_idx = end_idx + # [num_boxes x num_candidates] (index of candidate anchors) + candidate_idx = torch.cat(candidate_idx_list, dim=1) + + match_quality_matrix = self.similarity_fn(boxes, anchors) # [num_boxes x anchors] + candidate_ious = match_quality_matrix.gather(1, candidate_idx) # [num_boxes, n_candidates] + + # corner case, n_candidates<=1 will make iou_std_per_gt NaN + if candidate_idx.shape[1] <= 1: + matches = -1 * torch.ones((num_anchors,), dtype=torch.long, device=boxes.device) + matches[candidate_idx] = 0 + return match_quality_matrix, matches + + # compute adaptive iou threshold + iou_mean_per_gt = candidate_ious.mean(dim=1) # [num_boxes] + iou_std_per_gt = candidate_ious.std(dim=1) # [num_boxes] + iou_thresh_per_gt = iou_mean_per_gt + iou_std_per_gt # [num_boxes] + is_pos = candidate_ious >= iou_thresh_per_gt[:, None] # [num_boxes x n_candidates] + if self.debug: + print(f"Anchor matcher threshold: {iou_thresh_per_gt}") + + if self.center_in_gt: # can discard all candidates in case of very small objects :/ + # center point of selected anchors needs to lie within the ground truth + boxes_idx = ( + torch.arange(num_gt, device=boxes.device, dtype=torch.long)[:, None] + .expand_as(candidate_idx) + .contiguous() + ) # [num_boxes x n_candidates] + is_in_gt_ = centers_in_boxes( + anchors_center[candidate_idx.view(-1)], boxes[boxes_idx.view(-1)], eps=self.min_dist + ) + is_in_gt = convert_to_tensor(is_in_gt_) + is_pos = is_pos & is_in_gt.view_as(is_pos) # [num_boxes x n_candidates] + + # in case on anchor is assigned to multiple boxes, use box with highest IoU + # TODO: think about a better way to do this + for ng in range(num_gt): + candidate_idx[ng, :] += ng * num_anchors + ious_inf = torch.full_like(match_quality_matrix, -INF).view(-1) + index = candidate_idx.view(-1)[is_pos.view(-1)] + ious_inf[index] = match_quality_matrix.view(-1)[index] + ious_inf = ious_inf.view_as(match_quality_matrix) + + matched_vals, matches = ious_inf.to(COMPUTE_DTYPE).max(dim=0) + matches[matched_vals == -INF] = self.BELOW_LOW_THRESHOLD + # print(f"Num matches {(matches >= 0).sum()}, Adapt IoU {iou_thresh_per_gt}") + return match_quality_matrix, matches + + +MatcherType = TypeVar("MatcherType", bound=Matcher) diff --git a/monai/apps/detection/utils/__init__.py b/monai/apps/detection/utils/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/utils/__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/monai/apps/detection/utils/anchor_utils.py b/monai/apps/detection/utils/anchor_utils.py new file mode 100644 index 0000000000..c028228d95 --- /dev/null +++ b/monai/apps/detection/utils/anchor_utils.py @@ -0,0 +1,410 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE + +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +""" +This script is adapted from +https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py +""" + +from typing import List, Sequence, Union + +import torch +from torch import Tensor, nn + +from monai.utils import ensure_tuple +from monai.utils.misc import issequenceiterable +from monai.utils.module import look_up_option + + +class AnchorGenerator(nn.Module): + """ + This module is modified from torchvision to support both 2D and 3D images. + + Module that generates anchors for a set of feature maps and + image sizes. + + The module support computing anchors at multiple sizes and aspect ratios + per feature map. + + sizes and aspect_ratios should have the same number of elements, and it should + correspond to the number of feature maps. + + sizes[i] and aspect_ratios[i] can have an arbitrary number of elements. + For 2D images, anchor width and height w:h = 1:aspect_ratios[i,j] + For 3D images, anchor width, height, and depth w:h:d = 1:aspect_ratios[i,j,0]:aspect_ratios[i,j,1] + + AnchorGenerator will output a set of sizes[i] * aspect_ratios[i] anchors + per spatial location for feature map i. + + Args: + sizes: base size of each anchor. + len(sizes) is the number of feature maps, i.e., the number of output levels for + the feature pyramid network (FPN). + Each element of ``sizes`` is a Sequence which represents several anchor sizes for each feature map. + aspect_ratios: the aspect ratios of anchors. ``len(aspect_ratios) = len(sizes)``. + For 2D images, each element of ``aspect_ratios[i]`` is a Sequence of float. + For 3D images, each element of ``aspect_ratios[i]`` is a Sequence of 2 value Sequence. + indexing: choose from {``'ij'``, ``'xy'``}, optional, + Matrix (``'ij'``, default and recommended) or Cartesian (``'xy'``) indexing of output. + + - Matrix (``'ij'``, default and recommended) indexing keeps the original axis not changed. + - To use other monai detection components, please set ``indexing = 'ij'``. + - Cartesian (``'xy'``) indexing swaps axis 0 and 1. + - For 2D cases, monai ``AnchorGenerator(sizes, aspect_ratios, indexing='xy')`` and + ``torchvision.models.detection.anchor_utils.AnchorGenerator(sizes, aspect_ratios)`` are equivalent. + + + Reference:. + https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py + + Example: + .. code-block:: python + + # 2D example inputs for a 2-level feature maps + sizes = ((10,12,14,16), (20,24,28,32)) + base_aspect_ratios = (1., 0.5, 2.) + aspect_ratios = (base_aspect_ratios, base_aspect_ratios) + anchor_generator = AnchorGenerator(sizes, aspect_ratios) + + # 3D example inputs for a 2-level feature maps + sizes = ((10,12,14,16), (20,24,28,32)) + base_aspect_ratios = ((1., 1.), (1., 0.5), (0.5, 1.), (2., 2.)) + aspect_ratios = (base_aspect_ratios, base_aspect_ratios) + anchor_generator = AnchorGenerator(sizes, aspect_ratios) + """ + + __annotations__ = {"cell_anchors": List[torch.Tensor]} + + def __init__( + self, + sizes: Sequence[Sequence[int]] = ((20, 30, 40),), + aspect_ratios: Sequence = (((0.5, 1), (1, 0.5)),), + indexing: str = "ij", + ) -> None: + super().__init__() + + if not issequenceiterable(sizes[0]): + self.sizes = tuple((s,) for s in sizes) + else: + self.sizes = ensure_tuple(sizes) + if not issequenceiterable(aspect_ratios[0]): + aspect_ratios = (aspect_ratios,) * len(self.sizes) + + if len(self.sizes) != len(aspect_ratios): + raise ValueError( + "len(sizes) and len(aspect_ratios) should be equal. \ + It represents the number of feature maps." + ) + + spatial_dims = len(ensure_tuple(aspect_ratios[0][0])) + 1 + spatial_dims = look_up_option(spatial_dims, [2, 3]) + self.spatial_dims = spatial_dims + + self.indexing = look_up_option(indexing, ["ij", "xy"]) + + self.aspect_ratios = aspect_ratios + self.cell_anchors = [ + self.generate_anchors(size, aspect_ratio) for size, aspect_ratio in zip(self.sizes, aspect_ratios) + ] + + # This comment comes from torchvision. + # TODO: https://github.com/pytorch/pytorch/issues/26792 + # For every (aspect_ratios, scales) combination, output a zero-centered anchor with those values. + # (scales, aspect_ratios) are usually an element of zip(self.scales, self.aspect_ratios) + # This method assumes aspect ratio = height / width for an anchor. + def generate_anchors( + self, + scales: Sequence, + aspect_ratios: Sequence, + dtype: torch.dtype = torch.float32, + device: Union[torch.device, None] = None, + ) -> torch.Tensor: + """ + Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map. + + Args: + scales: a sequence which represents several anchor sizes for the current feature map. + aspect_ratios: a sequence which represents several aspect_ratios for the current feature map. + For 2D images, it is a Sequence of float aspect_ratios[j], + anchor width and height w:h = 1:aspect_ratios[j]. + For 3D images, it is a Sequence of 2 value Sequence aspect_ratios[j,0] and aspect_ratios[j,1], + anchor width, height, and depth w:h:d = 1:aspect_ratios[j,0]:aspect_ratios[j,1] + dtype: target data type of the output Tensor. + device: target device to put the output Tensor data. + + Returns: + For each s in scales, returns [s, s*aspect_ratios[j]] for 2D images, + and [s, s*aspect_ratios[j,0],s*aspect_ratios[j,1]] for 3D images. + """ + scales_t = torch.as_tensor(scales, dtype=dtype, device=device) # sized (N,) + aspect_ratios_t = torch.as_tensor(aspect_ratios, dtype=dtype, device=device) # sized (M,) or (M,2) + if (self.spatial_dims >= 3) and (len(aspect_ratios_t.shape) != 2): + ValueError( + f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \ + {len(aspect_ratios_t.shape)-1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}." + ) + + if (self.spatial_dims >= 3) and (aspect_ratios_t.shape[1] != self.spatial_dims - 1): + ValueError( + f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \ + shape (_,{self.spatial_dims-1}). But got aspect_ratios with shape {aspect_ratios_t.shape}." + ) + + # if 2d, w:h = 1:aspect_ratios + if self.spatial_dims == 2: + area_scale = torch.sqrt(aspect_ratios_t) + w_ratios = 1 / area_scale + h_ratios = area_scale + # if 3d, w:h:d = 1:aspect_ratios[:,0]:aspect_ratios[:,1] + elif self.spatial_dims == 3: + area_scale = torch.pow(aspect_ratios_t[:, 0] * aspect_ratios_t[:, 1], 1 / 3.0) + w_ratios = 1 / area_scale + h_ratios = aspect_ratios_t[:, 0] / area_scale + d_ratios = aspect_ratios_t[:, 1] / area_scale + + ws = (w_ratios[:, None] * scales_t[None, :]).view(-1) + hs = (h_ratios[:, None] * scales_t[None, :]).view(-1) + if self.spatial_dims == 2: + base_anchors = torch.stack([-ws, -hs, ws, hs], dim=1) / 2.0 + elif self.spatial_dims == 3: + ds = (d_ratios[:, None] * scales_t[None, :]).view(-1) + base_anchors = torch.stack([-ws, -hs, -ds, ws, hs, ds], dim=1) / 2.0 + + return base_anchors.round() + + def set_cell_anchors(self, dtype: torch.dtype, device: torch.device): + """ + Convert each element in self.cell_anchors to ``dtype`` and send to ``device``. + """ + self.cell_anchors = [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors] + + def num_anchors_per_location(self): + """ + Return number of anchor shapes for each feature map. + """ + return [c.shape[0] for c in self.cell_anchors] + + def grid_anchors(self, grid_sizes: List[List[int]], strides: List[List[Tensor]]) -> List[Tensor]: + """ + Every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:spatial_dims) + corresponds to a feature map. + It outputs g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a. + + Args: + grid_sizes: spatial size of the feature maps + strides: strides of the feature maps regarding to the original image + + Example: + .. code-block:: python + + grid_sizes = [[100,100],[50,50]] + strides = [[torch.tensor(2),torch.tensor(2)], [torch.tensor(4),torch.tensor(4)]] + """ + anchors = [] + cell_anchors = self.cell_anchors + if cell_anchors is None: + raise AssertionError + + if not (len(grid_sizes) == len(strides) == len(cell_anchors)): + raise ValueError( + "Anchors should be Tuple[Tuple[int]] because each feature " + "map could potentially have different sizes and aspect ratios. " + "There needs to be a match between the number of " + "feature maps passed and the number of sizes / aspect ratios specified." + ) + + for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors): + # for each feature map + device = base_anchors.device + + # compute anchor centers regarding to the image. + # shifts_centers is [x_center, y_center] or [x_center, y_center, z_center] + shifts_centers = [ + torch.arange(0, size[axis], dtype=torch.int32, device=device) * stride[axis] + for axis in range(self.spatial_dims) + ] + + # to support torchscript, cannot directly use torch.meshgrid(shifts_centers). + shifts_centers = list(torch.meshgrid(shifts_centers[: self.spatial_dims], indexing="ij")) + + for axis in range(self.spatial_dims): + # each element of shifts_centers is sized (HW,) or (HWD,) + shifts_centers[axis] = shifts_centers[axis].reshape(-1) + + # Expand to [x_center, y_center, x_center, y_center], + # or [x_center, y_center, z_center, x_center, y_center, z_center] + if self.indexing == "xy": + # Cartesian ('xy') indexing swaps axis 0 and 1. + shifts_centers[1], shifts_centers[0] = shifts_centers[0], shifts_centers[1] + shifts = torch.stack(shifts_centers * 2, dim=1) # sized (HW,4) or (HWD,6) + + # For every (base anchor, output anchor) pair, + # offset each zero-centered base anchor by the center of the output anchor. + anchors.append( + (shifts.view(-1, 1, self.spatial_dims * 2) + base_anchors.view(1, -1, self.spatial_dims * 2)).reshape( + -1, self.spatial_dims * 2 + ) # each element sized (AHWD,4) or (AHWD,6) + ) + + return anchors + + def forward(self, images: Tensor, feature_maps: List[Tensor]) -> List[Tensor]: + """ + Generate anchor boxes for each image. + + Args: + images: sized (B, C, W, H) or (B, C, W, H, D) + feature_maps: for FPN level i, feature_maps[i] is sized (B, C_i, W_i, H_i) or (B, C_i, W_i, H_i, D_i). + This input argument does not have to be the actual feature maps. + Any list variable with the same (C_i, W_i, H_i) or (C_i, W_i, H_i, D_i) as feature maps works. + + Return: + A list with length of B. Each element represents the anchors for this image. + The B elements are identical. + + Example: + .. code-block:: python + + images = torch.zeros((3,1,128,128,128)) + feature_maps = [torch.zeros((3,6,64,64,32)), torch.zeros((3,6,32,32,16))] + anchor_generator(images, feature_maps) + """ + grid_sizes = [list(feature_map.shape[-self.spatial_dims :]) for feature_map in feature_maps] + image_size = images.shape[-self.spatial_dims :] + batchsize = images.shape[0] + dtype, device = feature_maps[0].dtype, feature_maps[0].device + strides = [ + [ + torch.tensor(image_size[axis] // g[axis], dtype=torch.int64, device=device) + for axis in range(self.spatial_dims) + ] + for g in grid_sizes + ] + + self.set_cell_anchors(dtype, device) + anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides) + + anchors_per_image = torch.cat(list(anchors_over_all_feature_maps)) + return [anchors_per_image] * batchsize + + +class AnchorGeneratorWithAnchorShape(AnchorGenerator): + """ + Module that generates anchors for a set of feature maps and + image sizes, inherited from :py:class:`~monai.apps.detection.networks.utils.anchor_utils.AnchorGenerator` + + The module support computing anchors at multiple base anchor shapes + per feature map. + + ``feature_map_scales`` should have the same number of elements with the number of feature maps. + + base_anchor_shapes can have an arbitrary number of elements. + For 2D images, each element represents anchor width and height [w,h]. + For 2D images, each element represents anchor width, height, and depth [w,h,d]. + + AnchorGenerator will output a set of ``len(base_anchor_shapes)`` anchors + per spatial location for feature map ``i``. + + Args: + feature_map_scales: scale of anchors for each feature map, i.e., each output level of + the feature pyramid network (FPN). ``len(feature_map_scales)`` is the number of feature maps. + ``scale[i]*base_anchor_shapes`` represents the anchor shapes for feature map ``i``. + base_anchor_shapes: a sequence which represents several anchor shapes for one feature map. + For N-D images, it is a Sequence of N value Sequence. + indexing: choose from {'xy', 'ij'}, optional + Cartesian ('xy') or matrix ('ij', default) indexing of output. + Cartesian ('xy') indexing swaps axis 0 and 1, which is the setting inside torchvision. + matrix ('ij', default) indexing keeps the original axis not changed. + See also indexing in https://pytorch.org/docs/stable/generated/torch.meshgrid.html + + Example: + .. code-block:: python + + # 2D example inputs for a 2-level feature maps + feature_map_scales = (1, 2) + base_anchor_shapes = ((10, 10), (6, 12), (12, 6)) + anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes) + + # 3D example inputs for a 2-level feature maps + feature_map_scales = (1, 2) + base_anchor_shapes = ((10, 10, 10), (12, 12, 8), (10, 10, 6), (16, 16, 10)) + anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes) + """ + + __annotations__ = {"cell_anchors": List[torch.Tensor]} + + def __init__( + self, + feature_map_scales: Union[Sequence[int], Sequence[float]] = (1, 2, 4, 8), + base_anchor_shapes: Union[Sequence[Sequence[int]], Sequence[Sequence[float]]] = ( + (32, 32, 32), + (48, 20, 20), + (20, 48, 20), + (20, 20, 48), + ), + indexing: str = "ij", + ) -> None: + + nn.Module.__init__(self) + + spatial_dims = len(base_anchor_shapes[0]) + spatial_dims = look_up_option(spatial_dims, [2, 3]) + self.spatial_dims = spatial_dims + + self.indexing = look_up_option(indexing, ["ij", "xy"]) + + base_anchor_shapes_t = torch.Tensor(base_anchor_shapes) + self.cell_anchors = [self.generate_anchors_using_shape(s * base_anchor_shapes_t) for s in feature_map_scales] + + @staticmethod + def generate_anchors_using_shape( + anchor_shapes: torch.Tensor, dtype: torch.dtype = torch.float32, device: Union[torch.device, None] = None + ) -> torch.Tensor: + """ + Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map. + + Args: + anchor_shapes: [w, h] or [w, h, d], sized (N, spatial_dims), + represents N anchor shapes for the current feature map. + dtype: target data type of the output Tensor. + device: target device to put the output Tensor data. + + Returns: + For 2D images, returns [-w/2, -h/2, w/2, h/2]; + For 3D images, returns [-w/2, -h/2, -d/2, w/2, h/2, d/2] + """ + half_anchor_shapes = anchor_shapes / 2.0 + base_anchors = torch.cat([-half_anchor_shapes, half_anchor_shapes], dim=1) + return base_anchors.round().to(dtype=dtype, device=device) diff --git a/monai/apps/detection/utils/box_coder.py b/monai/apps/detection/utils/box_coder.py new file mode 100644 index 0000000000..036e23f242 --- /dev/null +++ b/monai/apps/detection/utils/box_coder.py @@ -0,0 +1,239 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/_utils.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE +# +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +This script is modified from torchvision to support N-D images, + +https://github.com/pytorch/vision/blob/main/torchvision/models/detection/_utils.py +""" + +import math +from typing import Sequence, Tuple, Union + +import torch +from torch import Tensor + +from monai.data.box_utils import COMPUTE_DTYPE, CenterSizeMode, StandardMode, convert_box_mode, is_valid_box_values +from monai.utils.module import look_up_option + + +def encode_boxes(gt_boxes: Tensor, proposals: Tensor, weights: Tensor) -> Tensor: + """ + Encode a set of proposals with respect to some reference ground truth (gt) boxes. + + Args: + gt_boxes: gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: boxes to be encoded, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + weights: the weights for ``(cx, cy, w, h) or (cx,cy,cz, w,h,d)`` + + Return: + encoded gt, target of box regression that is used to convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + + if gt_boxes.shape[0] != proposals.shape[0]: + raise ValueError("gt_boxes.shape[0] should be equal to proposals.shape[0].") + spatial_dims = look_up_option(len(weights), [4, 6]) // 2 + + if not is_valid_box_values(gt_boxes): + raise ValueError("gt_boxes is not valid. Please check if it contains empty boxes.") + if not is_valid_box_values(proposals): + raise ValueError("proposals is not valid. Please check if it contains empty boxes.") + + # implementation starts here + ex_cccwhd: Tensor = convert_box_mode(proposals, src_mode=StandardMode, dst_mode=CenterSizeMode) # type: ignore + gt_cccwhd: Tensor = convert_box_mode(gt_boxes, src_mode=StandardMode, dst_mode=CenterSizeMode) # type: ignore + targets_dxyz = ( + weights[:spatial_dims].unsqueeze(0) + * (gt_cccwhd[:, :spatial_dims] - ex_cccwhd[:, :spatial_dims]) + / ex_cccwhd[:, spatial_dims:] + ) + targets_dwhd = weights[spatial_dims:].unsqueeze(0) * torch.log( + gt_cccwhd[:, spatial_dims:] / ex_cccwhd[:, spatial_dims:] + ) + + targets = torch.cat((targets_dxyz, targets_dwhd), dim=1) + # torch.log may cause NaN or Inf + if torch.isnan(targets).any() or torch.isinf(targets).any(): + raise ValueError("targets is NaN or Inf.") + return targets + + +class BoxCoder: + """ + This class encodes and decodes a set of bounding boxes into + the representation used for training the regressors. + + Args: + weights: 4-element tuple or 6-element tuple + boxes_xform_clip: high threshold to prevent sending too large values into torch.exp() + + Example: + .. code-block:: python + + box_coder = BoxCoder(weights=[1., 1., 1., 1., 1., 1.]) + gt_boxes = torch.tensor([[1,2,1,4,5,6],[1,3,2,7,8,9]]) + proposals = gt_boxes + torch.rand(gt_boxes.shape) + rel_gt_boxes = box_coder.encode_single(gt_boxes, proposals) + gt_back = box_coder.decode_single(rel_gt_boxes, proposals) + # We expect gt_back to be equal to gt_boxes + """ + + def __init__(self, weights: Tuple[float], boxes_xform_clip: Union[float, None] = None) -> None: + if boxes_xform_clip is None: + boxes_xform_clip = math.log(1000.0 / 16) + self.spatial_dims = look_up_option(len(weights), [4, 6]) // 2 + self.weights = weights + self.boxes_xform_clip = boxes_xform_clip + + def encode(self, gt_boxes: Sequence[Tensor], proposals: Sequence[Tensor]) -> Tuple[Tensor]: + """ + Encode a set of proposals with respect to some ground truth (gt) boxes. + + Args: + gt_boxes: list of gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: list of boxes to be encoded, each element is Mx4 or Mx6 torch tensor. + The box mode is assumed to be ``StandardMode`` + + Return: + A tuple of encoded gt, target of box regression that is used to + convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + boxes_per_image = [len(b) for b in gt_boxes] + # concat the lists to do computation + concat_gt_boxes = torch.cat(tuple(gt_boxes), dim=0) + concat_proposals = torch.cat(tuple(proposals), dim=0) + concat_targets = self.encode_single(concat_gt_boxes, concat_proposals) + # split to tuple + targets: Tuple[Tensor] = concat_targets.split(boxes_per_image, 0) + return targets + + def encode_single(self, gt_boxes: Tensor, proposals: Tensor) -> Tensor: + """ + Encode proposals with respect to ground truth (gt) boxes. + + Args: + gt_boxes: gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: boxes to be encoded, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + + Return: + encoded gt, target of box regression that is used to convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + dtype = gt_boxes.dtype + device = gt_boxes.device + weights = torch.as_tensor(self.weights, dtype=dtype, device=device) + targets = encode_boxes(gt_boxes, proposals, weights) + return targets + + def decode(self, rel_codes: Tensor, reference_boxes: Sequence[Tensor]) -> Tensor: + """ + From a set of original reference_boxes and encoded relative box offsets, + + Args: + rel_codes: encoded boxes, Nx4 or Nx6 torch tensor. + boxes: a list of reference boxes, each element is Mx4 or Mx6 torch tensor. + The box mode is assumed to be ``StandardMode`` + + Return: + decoded boxes, Nx1x4 or Nx1x6 torch tensor. The box mode will be ``StandardMode`` + """ + if not isinstance(reference_boxes, Sequence) or (not isinstance(rel_codes, torch.Tensor)): + raise ValueError("Input arguments wrong type.") + boxes_per_image = [b.size(0) for b in reference_boxes] + # concat the lists to do computation + concat_boxes = torch.cat(tuple(reference_boxes), dim=0) + box_sum = 0 + for val in boxes_per_image: + box_sum += val + if box_sum > 0: + rel_codes = rel_codes.reshape(box_sum, -1) + pred_boxes = self.decode_single(rel_codes, concat_boxes) + if box_sum > 0: + pred_boxes = pred_boxes.reshape(box_sum, -1, 2 * self.spatial_dims) + return pred_boxes + + def decode_single(self, rel_codes: Tensor, reference_boxes: Tensor) -> Tensor: + """ + From a set of original boxes and encoded relative box offsets, + + Args: + rel_codes: encoded boxes, Nx(4*num_box_reg) or Nx(6*num_box_reg) torch tensor. + reference_boxes: reference boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + + Return: + decoded boxes, Nx(4*num_box_reg) or Nx(6*num_box_reg) torch tensor. The box mode will to be ``StandardMode`` + """ + reference_boxes = reference_boxes.to(rel_codes.dtype) + offset = reference_boxes.shape[-1] + + pred_boxes = [] + boxes_cccwhd = convert_box_mode(reference_boxes, src_mode=StandardMode, dst_mode=CenterSizeMode) + for axis in range(self.spatial_dims): + whd_axis = boxes_cccwhd[:, axis + self.spatial_dims] + ctr_xyz_axis = boxes_cccwhd[:, axis] + dxyz_axis = rel_codes[:, axis::offset] / self.weights[axis] + dwhd_axis = rel_codes[:, self.spatial_dims + axis :: offset] / self.weights[axis + self.spatial_dims] + # Prevent sending too large values into torch.exp() + dwhd_axis = torch.clamp(dwhd_axis.to(COMPUTE_DTYPE), max=self.boxes_xform_clip) + + pred_ctr_xyx_axis = dxyz_axis * whd_axis[:, None] + ctr_xyz_axis[:, None] + pred_whd_axis = torch.exp(dwhd_axis) * whd_axis[:, None] + pred_whd_axis = pred_whd_axis.to(dxyz_axis.dtype) + + # When convert float32 to float16, Inf or Nan may occur + if torch.isnan(pred_whd_axis).any() or torch.isinf(pred_whd_axis).any(): + raise ValueError("pred_whd_axis is NaN or Inf.") + + # Distance from center to box's corner. + c_to_c_whd_axis = ( + torch.tensor(0.5, dtype=pred_ctr_xyx_axis.dtype, device=pred_whd_axis.device) * pred_whd_axis + ) + + pred_boxes.append(pred_ctr_xyx_axis - c_to_c_whd_axis) + pred_boxes.append(pred_ctr_xyx_axis + c_to_c_whd_axis) + + pred_boxes = pred_boxes[::2] + pred_boxes[1::2] + pred_boxes_final = torch.stack(pred_boxes, dim=2).flatten(1) + return pred_boxes_final diff --git a/monai/apps/detection/utils/box_selector.py b/monai/apps/detection/utils/box_selector.py new file mode 100644 index 0000000000..de4de85ea0 --- /dev/null +++ b/monai/apps/detection/utils/box_selector.py @@ -0,0 +1,218 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE + +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +""" +Part of this script is adapted from +https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py +""" + +from typing import Callable, List, Tuple, Union + +import torch +from torch import Tensor + +from monai.data.box_utils import batched_nms, box_iou, clip_boxes_to_image +from monai.transforms.utils_pytorch_numpy_unification import floor_divide + + +class BoxSelector: + """ + Box selector which selects the predicted boxes. + The box selection is performed with the following steps: + + #. For each level, discard boxes with scores less than self.score_thresh. + #. For each level, keep boxes with top self.topk_candidates_per_level scores. + #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overlapping threshold nms_thresh. + #. For the whole image, keep boxes with top self.detections_per_img scores. + + Args: + apply_sigmoid: whether to apply sigmoid to get scores from classification logits + score_thresh: no box with scores less than score_thresh will be kept + topk_candidates_per_level: max number of boxes to keep for each level + nms_thresh: box overlapping threshold for NMS + detections_per_img: max number of boxes to keep for each image + + Example: + + .. code-block:: python + + input_param = { + "apply_sigmoid": True, + "score_thresh": 0.1, + "topk_candidates_per_level": 2, + "nms_thresh": 0.1, + "detections_per_img": 5, + } + box_selector = BoxSelector(**input_param) + boxes = [torch.randn([3,6]), torch.randn([7,6])] + logits = [torch.randn([3,3]), torch.randn([7,3])] + spatial_size = (8,8,8) + selected_boxes, selected_scores, selected_labels = box_selector.select_boxes_per_image( + boxes, logits, spatial_size + ) + """ + + def __init__( + self, + box_overlap_metric: Callable = box_iou, + apply_sigmoid: bool = True, + score_thresh: float = 0.05, + topk_candidates_per_level: int = 1000, + nms_thresh: float = 0.5, + detections_per_img: int = 300, + ): + self.box_overlap_metric = box_overlap_metric + + self.apply_sigmoid = apply_sigmoid + self.score_thresh = score_thresh + self.topk_candidates_per_level = topk_candidates_per_level + self.nms_thresh = nms_thresh + self.detections_per_img = detections_per_img + + def select_top_score_idx_per_level(self, logits: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Select indices with highest scores. + + The indice selection is performed with the following steps: + + #. If self.apply_sigmoid, get scores by applying sigmoid to logits. Otherwise, use logits as scores. + #. Discard indices with scores less than self.score_thresh + #. Keep indices with top self.topk_candidates_per_level scores + + Args: + logits: predicted classification logits, Tensor sized (N, num_classes) + + Return: + - topk_idxs: selected M indices, Tensor sized (M, ) + - selected_scores: selected M scores, Tensor sized (M, ) + - selected_labels: selected M labels, Tensor sized (M, ) + """ + num_classes = logits.shape[-1] + + # apply sigmoid to classification logits if asked + if self.apply_sigmoid: + scores = torch.sigmoid(logits.to(torch.float32)).flatten() + else: + scores = logits.flatten() + + # remove low scoring boxes + keep_idxs = scores > self.score_thresh + scores = scores[keep_idxs] + flatten_topk_idxs = torch.where(keep_idxs)[0] + + # keep only topk scoring predictions + num_topk = min(self.topk_candidates_per_level, flatten_topk_idxs.size(0)) + selected_scores, idxs = scores.to(torch.float32).topk( + num_topk + ) # half precision not implemented for cpu float16 + flatten_topk_idxs = flatten_topk_idxs[idxs] + + selected_labels = flatten_topk_idxs % num_classes + + topk_idxs = floor_divide(flatten_topk_idxs, num_classes) + return topk_idxs, selected_scores, selected_labels # type: ignore + + def select_boxes_per_image( + self, boxes_list: List[Tensor], logits_list: List[Tensor], spatial_size: Union[List[int], Tuple[int]] + ) -> Tuple[Tensor, Tensor, Tensor]: + """ + Postprocessing to generate detection result from classification logits and boxes. + + The box selection is performed with the following steps: + + #. For each level, discard boxes with scores less than self.score_thresh. + #. For each level, keep boxes with top self.topk_candidates_per_level scores. + #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overlapping threshold nms_thresh. + #. For the whole image, keep boxes with top self.detections_per_img scores. + + Args: + boxes_list: list of predicted boxes from a single image, + each element i is a Tensor sized (N_i, 2*spatial_dims) + logits_list: list of predicted classification logits from a single image, + each element i is a Tensor sized (N_i, num_classes) + spatial_size: spatial size of the image + + Return: + - selected boxes, Tensor sized (P, 2*spatial_dims) + - selected_scores, Tensor sized (P, ) + - selected_labels, Tensor sized (P, ) + """ + + if len(boxes_list) != len(logits_list): + raise ValueError( + "len(boxes_list) should equal to len(logits_list). " + f"Got len(boxes_list)={len(boxes_list)}, len(logits_list)={len(logits_list)}" + ) + + image_boxes = [] + image_scores = [] + image_labels = [] + + boxes_dtype = boxes_list[0].dtype + logits_dtype = logits_list[0].dtype + + for boxes_per_level, logits_per_level in zip(boxes_list, logits_list): + # select topk boxes for each level + topk_idxs: Tensor + topk_idxs, scores_per_level, labels_per_level = self.select_top_score_idx_per_level(logits_per_level) + boxes_per_level = boxes_per_level[topk_idxs] + + keep: Tensor + boxes_per_level, keep = clip_boxes_to_image( # type: ignore + boxes_per_level, spatial_size, remove_empty=True + ) + image_boxes.append(boxes_per_level) + image_scores.append(scores_per_level[keep]) + image_labels.append(labels_per_level[keep]) + + image_boxes_t: Tensor = torch.cat(image_boxes, dim=0) + image_scores_t: Tensor = torch.cat(image_scores, dim=0) + image_labels_t: Tensor = torch.cat(image_labels, dim=0) + + # non-maximum suppression on detected boxes from all levels + keep_t: Tensor = batched_nms( # type: ignore + image_boxes_t, + image_scores_t, + image_labels_t, + self.nms_thresh, + box_overlap_metric=self.box_overlap_metric, + max_proposals=self.detections_per_img, + ) + + selected_boxes = image_boxes_t[keep_t].to(boxes_dtype) + selected_scores = image_scores_t[keep_t].to(logits_dtype) + selected_labels = image_labels_t[keep_t] + + return selected_boxes, selected_scores, selected_labels diff --git a/monai/apps/detection/utils/detector_utils.py b/monai/apps/detection/utils/detector_utils.py new file mode 100644 index 0000000000..d7693da62c --- /dev/null +++ b/monai/apps/detection/utils/detector_utils.py @@ -0,0 +1,193 @@ +# 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 typing import Dict, List, Sequence, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import Tensor + +from monai.transforms.croppad.array import SpatialPad +from monai.transforms.utils import compute_divisible_spatial_size, convert_pad_mode +from monai.utils import PytorchPadMode, ensure_tuple_rep + + +def check_input_images(input_images: Union[List[Tensor], Tensor], spatial_dims: int) -> None: + """ + Validate the input dimensionality (raise a `ValueError` if invalid). + + Args: + input_images: It can be 1) a tensor sized (B, C, H, W) or (B, C, H, W, D), + or 2) a list of image tensors, each image i may have different size (C, H_i, W_i) or (C, H_i, W_i, D_i). + spatial_dims: number of spatial dimensions of the images, 2 or 3. + """ + if isinstance(input_images, Tensor): + if len(input_images.shape) != spatial_dims + 2: + raise ValueError( + "When input_images is a Tensor, its need to be (spatial_dims + 2)-D." + f"In this case, it should be a {(spatial_dims + 2)}-D Tensor, got Tensor shape {input_images.shape}." + ) + elif isinstance(input_images, List): + for img in input_images: + if len(img.shape) != spatial_dims + 1: + raise ValueError( + "When input_images is a List[Tensor], each element should have be (spatial_dims + 1)-D." + f"In this case, it should be a {(spatial_dims + 1)}-D Tensor, got Tensor shape {img.shape}." + ) + else: + raise ValueError("input_images needs to be a List[Tensor] or Tensor.") + return + + +def check_training_targets( + input_images: Union[List[Tensor], Tensor], + targets: Union[List[Dict[str, Tensor]], None], + spatial_dims: int, + target_label_key: str, + target_box_key: str, +) -> None: + """ + Validate the input images/targets during training (raise a `ValueError` if invalid). + + Args: + input_images: It can be 1) a tensor sized (B, C, H, W) or (B, C, H, W, D), + or 2) a list of image tensors, each image i may have different size (C, H_i, W_i) or (C, H_i, W_i, D_i). + targets: a list of dict. Each dict with two keys: target_box_key and target_label_key, + ground-truth boxes present in the image. + spatial_dims: number of spatial dimensions of the images, 2 or 3. + target_label_key: the expected key of target labels. + target_box_key: the expected key of target boxes. + """ + if targets is None: + raise ValueError("Please provide ground truth targets during training.") + + if len(input_images) != len(targets): + raise ValueError(f"len(input_images) should equal to len(targets), got {len(input_images)}, {len(targets)}.") + + for target in targets: + if (target_label_key not in target.keys()) or (target_box_key not in target.keys()): + raise ValueError( + f"{target_label_key} and {target_box_key} are expected keys in targets. Got {target.keys()}." + ) + + boxes = target[target_box_key] + if not isinstance(boxes, torch.Tensor): + raise ValueError(f"Expected target boxes to be of type Tensor, got {type(boxes)}.") + if len(boxes.shape) != 2 or boxes.shape[-1] != 2 * spatial_dims: + raise ValueError( + f"Expected target boxes to be a tensor " f"of shape [N, {2* spatial_dims}], got {boxes.shape}." + ) + return + + +def pad_images( + input_images: Union[List[Tensor], Tensor], + spatial_dims: int, + size_divisible: Union[int, Sequence[int]], + mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT, + **kwargs, +) -> Tuple[Tensor, List[List[int]]]: + """ + Pad the input images, so that the output spatial sizes are divisible by `size_divisible`. + It pads them at the end to create a (B, C, H, W) or (B, C, H, W, D) Tensor. + Padded size (H, W) or (H, W, D) is divisible by size_divisible. + Default padding uses constant padding with value 0.0 + + Args: + input_images: It can be 1) a tensor sized (B, C, H, W) or (B, C, H, W, D), + or 2) a list of image tensors, each image i may have different size (C, H_i, W_i) or (C, H_i, W_i, D_i). + spatial_dims: number of spatial dimensions of the images, 2D or 3D. + size_divisible: int or Sequence[int], is the expected pattern on the input image shape. + If an int, the same `size_divisible` will be applied to all the input spatial dimensions. + mode: available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + kwargs: other arguments for `torch.pad` function. + + Return: + - images, a (B, C, H, W) or (B, C, H, W, D) Tensor + - image_sizes, the original spatial size of each image + """ + size_divisible = ensure_tuple_rep(size_divisible, spatial_dims) + + # If input_images: Tensor + if isinstance(input_images, Tensor): + orig_size = list(input_images.shape[-spatial_dims:]) + new_size = compute_divisible_spatial_size(spatial_shape=orig_size, k=size_divisible) + all_pad_width = [(0, max(sp_i - orig_size[i], 0)) for i, sp_i in enumerate(new_size)] + pt_pad_width = [val for sublist in all_pad_width for val in sublist[::-1]][::-1] + if max(pt_pad_width) == 0: + # if there is no need to pad + return input_images, [orig_size] * input_images.shape[0] + mode_: str = convert_pad_mode(dst=input_images, mode=mode) + return F.pad(input_images, pt_pad_width, mode=mode_, **kwargs), [orig_size] * input_images.shape[0] + + # If input_images: List[Tensor]) + image_sizes = [img.shape[-spatial_dims:] for img in input_images] + in_channels = input_images[0].shape[0] + dtype = input_images[0].dtype + device = input_images[0].device + + # compute max_spatial_size + image_sizes_t = torch.tensor(image_sizes) + max_spatial_size_t, _ = torch.max(image_sizes_t, dim=0) + + if len(max_spatial_size_t) != spatial_dims or len(size_divisible) != spatial_dims: + raise ValueError(" Require len(max_spatial_size_t) == spatial_dims ==len(size_divisible).") + + max_spatial_size = compute_divisible_spatial_size(spatial_shape=list(max_spatial_size_t), k=size_divisible) + + # allocate memory for the padded images + images = torch.zeros([len(image_sizes), in_channels] + max_spatial_size, dtype=dtype, device=device) + + # Use `SpatialPad` to match sizes, padding in the end will not affect boxes + padder = SpatialPad(spatial_size=max_spatial_size, method="end", mode=mode, **kwargs) + for idx, img in enumerate(input_images): + images[idx, ...] = padder(img) + + return images, [list(ss) for ss in image_sizes] + + +def preprocess_images( + input_images: Union[List[Tensor], Tensor], + spatial_dims: int, + size_divisible: Union[int, Sequence[int]], + mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT, + **kwargs, +) -> Tuple[Tensor, List[List[int]]]: + """ + Preprocess the input images, including + + - validate of the inputs + - pad the inputs so that the output spatial sizes are divisible by `size_divisible`. + It pads them at the end to create a (B, C, H, W) or (B, C, H, W, D) Tensor. + Padded size (H, W) or (H, W, D) is divisible by size_divisible. + Default padding uses constant padding with value 0.0 + + Args: + input_images: It can be 1) a tensor sized (B, C, H, W) or (B, C, H, W, D), + or 2) a list of image tensors, each image i may have different size (C, H_i, W_i) or (C, H_i, W_i, D_i). + spatial_dims: number of spatial dimensions of the images, 2 or 3. + size_divisible: int or Sequence[int], is the expected pattern on the input image shape. + If an int, the same `size_divisible` will be applied to all the input spatial dimensions. + mode: available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + kwargs: other arguments for `torch.pad` function. + + Return: + - images, a (B, C, H, W) or (B, C, H, W, D) Tensor + - image_sizes, the original spatial size of each image + """ + check_input_images(input_images, spatial_dims) + size_divisible = ensure_tuple_rep(size_divisible, spatial_dims) + + return pad_images(input_images, spatial_dims, size_divisible, mode, **kwargs) diff --git a/monai/apps/detection/utils/hard_negative_sampler.py b/monai/apps/detection/utils/hard_negative_sampler.py new file mode 100644 index 0000000000..3067a41a0e --- /dev/null +++ b/monai/apps/detection/utils/hard_negative_sampler.py @@ -0,0 +1,306 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/sampler.py +# which has the following license... +# https://github.com/MIC-DKFZ/nnDetection/blob/main/LICENSE +# +# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany +# 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. + +""" +The functions in this script are adapted from nnDetection, +https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/sampler.py +""" + +import logging +from abc import ABC +from typing import List, Tuple + +import torch +from torch import Tensor + + +class HardNegativeSamplerBase(ABC): + """ + Base class of hard negative sampler. + + Hard negative sampler is used to suppress false positive rate in classification tasks. + During training, it select negative samples with high prediction scores. + + The training workflow is described as the follows: + 1) forward network and get prediction scores (classification prob/logits) for all the samples; + 2) use hard negative sampler to choose negative samples with high prediction scores and some positive samples; + 3) compute classification loss for the selected samples; + 4) do back propagation. + + Args: + pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from + ``num_neg * pool_size`` negative samples with the highest prediction scores. + Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', + i.e., negative samples with lower prediction scores. + """ + + def __init__(self, pool_size: float = 10) -> None: + self.pool_size = pool_size + + def select_negatives(self, negative: Tensor, num_neg: int, fg_probs: Tensor) -> Tensor: + """ + Select hard negative samples. + + Args: + negative: indices of all the negative samples, sized (P,), + where P is the number of negative samples + num_neg: number of negative samples to sample + fg_probs: maximum foreground prediction scores (probability) across all the classes + for each sample, sized (A,), where A is the the number of samples. + + Returns: + binary mask of negative samples to choose, sized (A,), + where A is the the number of samples in one image + """ + if negative.numel() > fg_probs.numel(): + raise ValueError("The number of negative samples should not be larger than the number of all samples.") + + # sample pool size is ``num_neg * self.pool_size`` + pool = int(num_neg * self.pool_size) + pool = min(negative.numel(), pool) # protect against not enough negatives + + # create a sample pool of highest scoring negative samples + _, negative_idx_pool = fg_probs[negative].to(torch.float32).topk(pool, dim=0, sorted=True) + hard_negative = negative[negative_idx_pool] + + # select negatives from pool + perm2 = torch.randperm(hard_negative.numel(), device=hard_negative.device)[:num_neg] + selected_neg_idx = hard_negative[perm2] + + # output a binary mask with same size of fg_probs that indicates selected negative samples. + neg_mask = torch.zeros_like(fg_probs, dtype=torch.uint8) + neg_mask[selected_neg_idx] = 1 + return neg_mask + + +class HardNegativeSampler(HardNegativeSamplerBase): + """ + HardNegativeSampler is used to suppress false positive rate in classification tasks. + During training, it selects negative samples with high prediction scores. + + The training workflow is described as the follows: + 1) forward network and get prediction scores (classification prob/logits) for all the samples; + 2) use hard negative sampler to choose negative samples with high prediction scores and some positive samples; + 3) compute classification loss for the selected samples; + 4) do back propagation. + + Args: + batch_size_per_image: number of training samples to be randomly selected per image + positive_fraction: percentage of positive elements in the selected samples + min_neg: minimum number of negative samples to select if possible. + pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from + ``num_neg * pool_size`` negative samples with the highest prediction scores. + Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', + i.e., negative samples with lower prediction scores. + """ + + def __init__( + self, batch_size_per_image: int, positive_fraction: float, min_neg: int = 1, pool_size: float = 10 + ) -> None: + super().__init__(pool_size=pool_size) + self.min_neg = min_neg + self.batch_size_per_image = batch_size_per_image + self.positive_fraction = positive_fraction + logging.info("Sampling hard negatives on a per batch basis") + + def __call__(self, target_labels: List[Tensor], concat_fg_probs: Tensor) -> Tuple[List[Tensor], List[Tensor]]: + """ + Select positives and hard negatives from list samples per image. + Hard negative sampler will be applied to each image independently. + + Args: + target_labels: list of labels per image. + For image i in the batch, target_labels[i] is a Tensor sized (A_i,), + where A_i is the number of samples in image i. + Positive samples have positive labels, negative samples have label 0. + concat_fg_probs: concatenated maximum foreground probability for all the images, sized (R,), + where R is the sum of all samples inside one batch, i.e., R = A_0 + A_1 + ... + + Returns: + - list of binary mask for positive samples + - list of binary mask for negative samples + + Example: + .. code-block:: python + + sampler = HardNegativeSampler( + batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 + ) + # two images with different number of samples + target_labels = [ torch.tensor([0,1]), torch.tensor([1,0,2,1])] + concat_fg_probs = torch.rand(6) + pos_idx_list, neg_idx_list = sampler(target_labels, concat_fg_probs) + """ + samples_per_image = [samples_in_image.shape[0] for samples_in_image in target_labels] + fg_probs = concat_fg_probs.split(samples_per_image, 0) + return self.select_samples_img_list(target_labels, fg_probs) + + def select_samples_img_list( + self, target_labels: List[Tensor], fg_probs: List[Tensor] + ) -> Tuple[List[Tensor], List[Tensor]]: + """ + Select positives and hard negatives from list samples per image. + Hard negative sampler will be applied to each image independently. + + Args: + target_labels: list of labels per image. + For image i in the batch, target_labels[i] is a Tensor sized (A_i,), + where A_i is the number of samples in image i. + Positive samples have positive labels, negative samples have label 0. + fg_probs: list of maximum foreground probability per images, + For image i in the batch, target_labels[i] is a Tensor sized (A_i,), + where A_i is the number of samples in image i. + + Returns: + - list of binary mask for positive samples + - list binary mask for negative samples + + Example: + .. code-block:: python + + sampler = HardNegativeSampler( + batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 + ) + # two images with different number of samples + target_labels = [ torch.tensor([0,1]), torch.tensor([1,0,2,1])] + fg_probs = [ torch.rand(2), torch.rand(4)] + pos_idx_list, neg_idx_list = sampler.select_samples_img_list(target_labels, fg_probs) + """ + pos_idx = [] + neg_idx = [] + + if len(target_labels) != len(fg_probs): + raise ValueError( + "Require len(target_labels) == len(fg_probs). " + f"Got len(target_labels)={len(target_labels)}, len(fg_probs)={len(fg_probs)}." + ) + for labels_per_img, fg_probs_per_img in zip(target_labels, fg_probs): + pos_idx_per_image_mask, neg_idx_per_image_mask = self.select_samples_per_img( + labels_per_img, fg_probs_per_img + ) + pos_idx.append(pos_idx_per_image_mask) + neg_idx.append(neg_idx_per_image_mask) + + return pos_idx, neg_idx + + def select_samples_per_img(self, labels_per_img: Tensor, fg_probs_per_img: Tensor) -> Tuple[Tensor, Tensor]: + """ + Select positives and hard negatives from samples. + + Args: + labels_per_img: labels, sized (A,). + Positive samples have positive labels, negative samples have label 0. + fg_probs_per_img: maximum foreground probability, sized (A,) + + Returns: + - binary mask for positive samples, sized (A,) + - binary mask for negative samples, sized (A,) + + Example: + .. code-block:: python + + sampler = HardNegativeSampler( + batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 + ) + # two images with different number of samples + target_labels = torch.tensor([1,0,2,1]) + fg_probs = torch.rand(4) + pos_idx, neg_idx = sampler.select_samples_per_img(target_labels, fg_probs) + """ + # for each image, find positive sample indices and negative sample indices + if labels_per_img.numel() != fg_probs_per_img.numel(): + raise ValueError("labels_per_img and fg_probs_per_img should have same number of elements.") + + positive = torch.where(labels_per_img >= 1)[0] + negative = torch.where(labels_per_img == 0)[0] + + num_pos = self.get_num_pos(positive) + pos_idx_per_image_mask = self.select_positives(positive, num_pos, labels_per_img) + + num_neg = self.get_num_neg(negative, num_pos) + neg_idx_per_image_mask = self.select_negatives(negative, num_neg, fg_probs_per_img) + + return pos_idx_per_image_mask, neg_idx_per_image_mask + + def get_num_pos(self, positive: torch.Tensor) -> int: + """ + Number of positive samples to draw + + Args: + positive: indices of positive samples + + Returns: + number of positive sample + """ + # positive sample sampling + num_pos = int(self.batch_size_per_image * self.positive_fraction) + # protect against not enough positive examples + num_pos = min(positive.numel(), num_pos) + return num_pos + + def get_num_neg(self, negative: torch.Tensor, num_pos: int) -> int: + """ + Sample enough negatives to fill up ``self.batch_size_per_image`` + + Args: + negative: indices of positive samples + num_pos: number of positive samples to draw + + Returns: + number of negative samples + """ + # always assume at least one pos sample was sampled + num_neg = int(max(1, num_pos) * abs(1 - 1.0 / float(self.positive_fraction))) + # protect against not enough negative examples and sample at least self.min_neg if possible + num_neg = min(negative.numel(), max(num_neg, self.min_neg)) + return num_neg + + def select_positives(self, positive: Tensor, num_pos: int, labels: Tensor) -> Tensor: + """ + Select positive samples + + Args: + positive: indices of positive samples, sized (P,), + where P is the number of positive samples + num_pos: number of positive samples to sample + labels: labels for all samples, sized (A,), + where A is the number of samples. + + Returns: + binary mask of positive samples to choose, sized (A,), + where A is the number of samples in one image + """ + if positive.numel() > labels.numel(): + raise ValueError("The number of positive samples should not be larger than the number of all samples.") + + perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos] + pos_idx_per_image = positive[perm1] + + # output a binary mask with same size of labels that indicates selected positive samples. + pos_idx_per_image_mask = torch.zeros_like(labels, dtype=torch.uint8) + pos_idx_per_image_mask[pos_idx_per_image] = 1 + return pos_idx_per_image_mask diff --git a/monai/apps/detection/utils/predict_utils.py b/monai/apps/detection/utils/predict_utils.py new file mode 100644 index 0000000000..a11aa97ce7 --- /dev/null +++ b/monai/apps/detection/utils/predict_utils.py @@ -0,0 +1,135 @@ +# 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 typing import Dict, List, Optional + +import torch +from torch import Tensor + +from monai.inferers import SlidingWindowInferer + + +def ensure_dict_value_to_list_(head_outputs: Dict[str, List[Tensor]], keys: Optional[List[str]] = None) -> None: + """ + An in-place function. We expect ``head_outputs`` to be Dict[str, List[Tensor]]. + Yet if it is Dict[str, Tensor], this func converts it to Dict[str, List[Tensor]]. + It will be modified in-place. + + Args: + head_outputs: a Dict[str, List[Tensor]] or Dict[str, Tensor], will be modifier in-place + keys: the keys in head_output that need to have value type List[Tensor]. If not provided, will use head_outputs.keys(). + """ + if keys is None: + keys = list(head_outputs.keys()) + + for k in keys: + value_k = head_outputs[k] # Tensor or List[Tensor] + # convert value_k to List[Tensor] + if isinstance(value_k, Tensor): + head_outputs[k] = [value_k] + elif isinstance(value_k[0], Tensor): + head_outputs[k] = list(value_k) + else: + raise ValueError("The output of network should be Dict[str, List[Tensor]] or Dict[str, Tensor].") + + +def check_dict_values_same_length(head_outputs: Dict[str, List[Tensor]], keys: Optional[List[str]] = None) -> None: + """ + We expect the values in ``head_outputs``: Dict[str, List[Tensor]] to have the same length. + Will raise ValueError if not. + + Args: + head_outputs: a Dict[str, List[Tensor]] or Dict[str, Tensor] + keys: the keys in head_output that need to have values (List) with same length. + If not provided, will use head_outputs.keys(). + """ + if keys is None: + keys = list(head_outputs.keys()) + + num_output_levels_list: List[int] = [len(head_outputs[k]) for k in keys] + num_output_levels = torch.unique(torch.tensor(num_output_levels_list)) + if len(num_output_levels) != 1: + raise ValueError(f"The values in the input dict should have the same length, Got {num_output_levels_list}.") + + +def _network_sequence_output(images: Tensor, network, keys: Optional[List[str]] = None) -> List[Tensor]: + """ + Decompose the output of network (a dict) into a list. + + Args: + images: input of the network + keys: the keys in the network output whose values will be output in this func. + If not provided, will use all keys. + + Return: + network output values concat to a single List[Tensor] + """ + head_outputs = network(images) + ensure_dict_value_to_list_(head_outputs, keys) + if keys is None: + keys = list(head_outputs.keys()) + check_dict_values_same_length(head_outputs, keys) + head_outputs_sequence = [] + for k in keys: + head_outputs_sequence += list(head_outputs[k]) + return head_outputs_sequence + + +def predict_with_inferer( + images: Tensor, network, keys: List[str], inferer: Optional[SlidingWindowInferer] = None +) -> Dict[str, List[Tensor]]: + """ + Predict network dict output with an inferer. Compared with directly output network(images), + it enables a sliding window inferer that can be used to handle large inputs. + + Args: + images: input of the network, Tensor sized (B, C, H, W) or (B, C, H, W, D) + network: a network that takes an image Tensor sized (B, C, H, W) or (B, C, H, W, D) as input + and outputs a dictionary Dict[str, List[Tensor]] or Dict[str, Tensor]. + keys: the keys in the output dict, should be network output keys or a subset of them. + inferer: a SlidingWindowInferer to handle large inputs. + + Return: + The predicted head_output from network, a Dict[str, List[Tensor]] + + Example: + .. code-block:: python + + # define a naive network + import torch + import monai + class NaiveNet(torch.nn.Module): + def __init__(self, ): + super().__init__() + + def forward(self, images: torch.Tensor): + return {"cls": torch.randn(images.shape), "box_reg": [torch.randn(images.shape)]} + + # create a predictor + network = NaiveNet() + inferer = monai.inferers.SlidingWindowInferer( + roi_size = (128, 128, 128), + overlap = 0.25, + cache_roi_weight_map = True, + ) + network_output_keys=["cls", "box_reg"] + images = torch.randn((2, 3, 512, 512, 512)) # a large input + head_outputs = predict_with_inferer(images, network, network_output_keys, inferer) + + """ + if inferer is None: + raise ValueError("Please set inferer as a monai.inferers.inferer.SlidingWindowInferer(*)") + head_outputs_sequence = inferer(images, _network_sequence_output, network, keys=keys) + num_output_levels: int = len(head_outputs_sequence) // len(keys) + head_outputs = {} + for i, k in enumerate(keys): + head_outputs[k] = list(head_outputs_sequence[num_output_levels * i : num_output_levels * (i + 1)]) + return head_outputs diff --git a/monai/apps/nuclick/__init__.py b/monai/apps/nuclick/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/nuclick/__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/monai/apps/nuclick/transforms.py b/monai/apps/nuclick/transforms.py new file mode 100644 index 0000000000..28c6417a42 --- /dev/null +++ b/monai/apps/nuclick/transforms.py @@ -0,0 +1,532 @@ +# 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. + +import math +import random +from typing import Any, Tuple, Union + +import numpy as np + +from monai.config import KeysCollection +from monai.transforms import MapTransform, Randomizable, SpatialPad +from monai.utils import StrEnum, optional_import + +measure, _ = optional_import("skimage.measure") +morphology, _ = optional_import("skimage.morphology") + + +class NuclickKeys(StrEnum): + """ + Keys for nuclick transforms. + """ + + IMAGE = "image" + LABEL = "label" + OTHERS = "others" # key of other labels from the binary mask which are not being used for training + FOREGROUND = "foreground" + + CENTROID = "centroid" # key where the centroid values are stored + MASK_VALUE = "mask_value" + LOCATION = "location" + + NUC_POINTS = "nuc_points" + BOUNDING_BOXES = "bounding_boxes" + IMG_HEIGHT = "img_height" + IMG_WIDTH = "img_width" + + +class FlattenLabeld(MapTransform): + """ + FlattenLabeld creates labels per closed object contour (defined by a connectivity). For e.g if there are + 12 small regions of 1's it will delineate them into 12 different label classes + + Args: + connectivity: Max no. of orthogonal hops to consider a pixel/voxel as a neighbor. Refer skimage.measure.label + allow_missing_keys: don't raise exception if key is missing. + """ + + def __init__(self, keys: KeysCollection, connectivity: int = 1, allow_missing_keys: bool = False): + super().__init__(keys, allow_missing_keys) + self.connectivity = connectivity + + def __call__(self, data): + d = dict(data) + for key in self.keys: + d[key] = measure.label(d[key], connectivity=self.connectivity).astype(np.uint8) + return d + + +class ExtractPatchd(MapTransform): + """ + Extracts a patch from the given image and label, however it is based on the centroid location. + The centroid location is a 2D coordinate (H, W). The extracted patch is extracted around the centroid, + if the centroid is towards the edge, the centroid will not be the center of the image as the patch will be + extracted from the edges onwards + + Args: + keys: image, label + centroid_key: key where the centroid values are stored, defaults to ``"centroid"`` + patch_size: size of the extracted patch + allow_missing_keys: don't raise exception if key is missing. + pad_kwargs: other arguments for the SpatialPad transform + """ + + def __init__( + self, + keys: KeysCollection, + centroid_key: str = NuclickKeys.CENTROID, + patch_size: Union[Tuple[int, int], int] = 128, + allow_missing_keys: bool = False, + **kwargs: Any, + ): + super().__init__(keys, allow_missing_keys) + self.centroid_key = centroid_key + self.patch_size = patch_size + self.kwargs = kwargs + + def __call__(self, data): + d = dict(data) + + centroid = d[self.centroid_key] # create mask based on centroid (select nuclei based on centroid) + roi_size = (self.patch_size, self.patch_size) + + for key in self.keys: + img = d[key] + x_start, x_end, y_start, y_end = self.bbox(self.patch_size, centroid, img.shape[-2:]) + cropped = img[:, x_start:x_end, y_start:y_end] + d[key] = SpatialPad(spatial_size=roi_size, **self.kwargs)(cropped) + return d + + def bbox(self, patch_size, centroid, size): + x, y = centroid + m, n = size + + x_start = int(max(x - patch_size / 2, 0)) + y_start = int(max(y - patch_size / 2, 0)) + x_end = x_start + patch_size + y_end = y_start + patch_size + if x_end > m: + x_end = m + x_start = m - patch_size + if y_end > n: + y_end = n + y_start = n - patch_size + return x_start, x_end, y_start, y_end + + +class SplitLabeld(MapTransform): + """ + Extracts a single label from all the given classes, the single label is defined by mask_value, the remaining + labels are kept in others + + Args: + label: key of the label source + others: other labels storage key, defaults to ``"others"`` + mask_value: the mask_value that will be kept for binarization of the label, defaults to ``"mask_value"`` + min_area: The smallest allowable object size. + """ + + def __init__( + self, + keys: KeysCollection, + # label: str = NuclickKeys.LABEL, + others: str = NuclickKeys.OTHERS, + mask_value: str = NuclickKeys.MASK_VALUE, + min_area: int = 5, + ): + + # self.label = label + super().__init__(keys, allow_missing_keys=False) + self.others = others + self.mask_value = mask_value + self.min_area = min_area + + def __call__(self, data): + d = dict(data) + + if len(self.keys) > 1: + print("Only 'label' key is supported, more than 1 key was found") + return None + + for key in self.keys: + self.label = key + + label = d[self.label] + mask_value = d[self.mask_value] + mask = np.uint8(label == mask_value) + others = (1 - mask) * label + others = self._mask_relabeling(others[0], min_area=self.min_area)[np.newaxis] + d[self.label] = mask + d[self.others] = others + return d + + def _mask_relabeling(self, mask, min_area=5): + res = np.zeros_like(mask) + for l in np.unique(mask): + if l == 0: + continue + + m = measure.label(mask == l, connectivity=1) + for stat in measure.regionprops(m): + if stat.area > min_area: + res[stat.coords[:, 0], stat.coords[:, 1]] = l + return res + + +class FilterImaged(MapTransform): + """ + Filters Green and Gray channel of the image using an allowable object size, this pre-processing transform + is specific towards NuClick training process. More details can be referred in this paper Koohbanani, + Navid Alemi, et al. "NuClick: a deep learning framework for interactive segmentation of microscopic images." + Medical Image Analysis 65 (2020): 101771. + + Args: + min_size: The smallest allowable object size + allow_missing_keys: don't raise exception if key is missing. + """ + + def __init__(self, keys: KeysCollection, min_size: int = 500, allow_missing_keys: bool = False): + super().__init__(keys, allow_missing_keys) + self.min_size = min_size + + def __call__(self, data): + d = dict(data) + for key in self.keys: + img = d[key] + d[key] = self.filter(img) + return d + + def filter(self, rgb): + mask_not_green = self.filter_green_channel(rgb) + mask_not_gray = self.filter_grays(rgb) + mask_gray_green = mask_not_gray & mask_not_green + mask = ( + self.filter_remove_small_objects(mask_gray_green, min_size=self.min_size) + if self.min_size + else mask_gray_green + ) + + return rgb * np.dstack([mask, mask, mask]) + + def filter_green_channel( + self, img_np, green_thresh=200, avoid_overmask=True, overmask_thresh=90, output_type="bool" + ): + g = img_np[:, :, 1] + gr_ch_mask = (g < green_thresh) & (g > 0) + mask_percentage = self.mask_percent(gr_ch_mask) + if (mask_percentage >= overmask_thresh) and (green_thresh < 255) and (avoid_overmask is True): + new_green_thresh = math.ceil((255 - green_thresh) / 2 + green_thresh) + gr_ch_mask = self.filter_green_channel( + img_np, new_green_thresh, avoid_overmask, overmask_thresh, output_type + ) + return gr_ch_mask + + def filter_grays(self, rgb, tolerance=15): + rg_diff = abs(rgb[:, :, 0] - rgb[:, :, 1]) <= tolerance + rb_diff = abs(rgb[:, :, 0] - rgb[:, :, 2]) <= tolerance + gb_diff = abs(rgb[:, :, 1] - rgb[:, :, 2]) <= tolerance + return ~(rg_diff & rb_diff & gb_diff) + + def mask_percent(self, img_np): + if (len(img_np.shape) == 3) and (img_np.shape[2] == 3): + np_sum = img_np[:, :, 0] + img_np[:, :, 1] + img_np[:, :, 2] + mask_percentage = 100 - np.count_nonzero(np_sum) / np_sum.size * 100 + else: + mask_percentage = 100 - np.count_nonzero(img_np) / img_np.size * 100 + return mask_percentage + + def filter_remove_small_objects(self, img_np, min_size=3000, avoid_overmask=True, overmask_thresh=95): + rem_sm = morphology.remove_small_objects(img_np.astype(bool), min_size=min_size) + mask_percentage = self.mask_percent(rem_sm) + if (mask_percentage >= overmask_thresh) and (min_size >= 1) and (avoid_overmask is True): + new_min_size = round(min_size / 2) + rem_sm = self.filter_remove_small_objects(img_np, new_min_size, avoid_overmask, overmask_thresh) + return rem_sm + + +class AddPointGuidanceSignald(Randomizable, MapTransform): + """ + Adds Guidance Signal to the input image + + Args: + image: key of source image, defaults to ``"image"`` + label: key of source label, defaults to ``"label"`` + others: source others (other labels from the binary mask which are not being used for training) + defaults to ``"others"`` + drop_rate: probability of dropping the signal, defaults to ``0.5`` + jitter_range: noise added to the points in the point mask for exclusion mask, defaults to ``3`` + """ + + def __init__( + self, + image: str = NuclickKeys.IMAGE, + label: str = NuclickKeys.LABEL, + others: str = NuclickKeys.OTHERS, + drop_rate: float = 0.5, + jitter_range: int = 3, + ): + MapTransform.__init__(self, image) + + self.image = image + self.label = label + self.others = others + self.drop_rate = drop_rate + self.jitter_range = jitter_range + + def __call__(self, data): + d = dict(data) + + image = d[self.image] + mask = d[self.label] + others = d[self.others] + + inc_sig = self.inclusion_map(mask[0]) + exc_sig = self.exclusion_map(others[0], drop_rate=self.drop_rate, jitter_range=self.jitter_range) + + image = np.concatenate((image, inc_sig[np.newaxis, ...], exc_sig[np.newaxis, ...]), axis=0) + d[self.image] = image + return d + + def inclusion_map(self, mask): + point_mask = np.zeros_like(mask) + indices = np.argwhere(mask > 0) + if len(indices) > 0: + idx = np.random.randint(0, len(indices)) + point_mask[indices[idx, 0], indices[idx, 1]] = 1 + + return point_mask + + def exclusion_map(self, others, jitter_range=3, drop_rate=0.5): + point_mask = np.zeros_like(others) + if drop_rate == 1.0: + return point_mask + + max_x = point_mask.shape[0] - 1 + max_y = point_mask.shape[1] - 1 + stats = measure.regionprops(others) + for stat in stats: + x, y = stat.centroid + if np.random.choice([True, False], p=[drop_rate, 1 - drop_rate]): + continue + + # random jitter + x = int(math.floor(x)) + random.randint(a=-jitter_range, b=jitter_range) + y = int(math.floor(y)) + random.randint(a=-jitter_range, b=jitter_range) + x = min(max(0, x), max_x) + y = min(max(0, y), max_y) + point_mask[x, y] = 1 + + return point_mask + + +class AddClickSignalsd(MapTransform): + """ + Adds Click Signal to the input image + + Args: + image: source image, defaults to ``"image"`` + foreground: 2D click indices as list, defaults to ``"foreground"`` + bb_size: single integer size, defines a bounding box like (bb_size, bb_size) + """ + + def __init__(self, image: str = NuclickKeys.IMAGE, foreground: str = NuclickKeys.FOREGROUND, bb_size: int = 128): + self.image = image + self.foreground = foreground + self.bb_size = bb_size + + def __call__(self, data): + d = dict(data) + + location = d.get(NuclickKeys.LOCATION.value, (0, 0)) + tx, ty = location[0], location[1] + pos = d.get(self.foreground) + pos = (np.array(pos) - (tx, ty)).astype(int).tolist() if pos else [] + + cx = [xy[0] for xy in pos] + cy = [xy[1] for xy in pos] + + img = d[self.image].astype(np.uint8) + img_width = img.shape[-1] + img_height = img.shape[-2] + + click_map, bounding_boxes = self.get_clickmap_boundingbox( + cx=cx, cy=cy, m=img_height, n=img_width, bb=self.bb_size + ) + + patches, nuc_points, other_points = self.get_patches_and_signals( + img=img, + click_map=click_map, + bounding_boxes=bounding_boxes, + cx=cx, + cy=cy, + m=img_height, + n=img_width, + bb=self.bb_size, + ) + patches = patches / 255 + + d[NuclickKeys.BOUNDING_BOXES.value] = bounding_boxes + d[NuclickKeys.IMG_WIDTH.value] = img_width + d[NuclickKeys.IMG_HEIGHT.value] = img_height + d[NuclickKeys.NUC_POINTS.value] = nuc_points + + d[self.image] = np.concatenate((patches, nuc_points, other_points), axis=1).astype(dtype=np.float32) + return d + + def get_clickmap_boundingbox(self, cx, cy, m, n, bb=128): + click_map = np.zeros((m, n), dtype=np.uint8) + + x_del_indices = {i for i in range(len(cx)) if cx[i] >= n or cx[i] < 0} + y_del_indices = {i for i in range(len(cy)) if cy[i] >= m or cy[i] < 0} + del_indices = list(x_del_indices.union(y_del_indices)) + cx = np.delete(cx, del_indices) + cy = np.delete(cy, del_indices) + + click_map[cy, cx] = 1 + bounding_boxes = [] + for i in range(len(cx)): + x_start = cx[i] - bb // 2 + y_start = cy[i] - bb // 2 + if x_start < 0: + x_start = 0 + if y_start < 0: + y_start = 0 + x_end = x_start + bb - 1 + y_end = y_start + bb - 1 + if x_end > n - 1: + x_end = n - 1 + x_start = x_end - bb + 1 + if y_end > m - 1: + y_end = m - 1 + y_start = y_end - bb + 1 + bounding_boxes.append([x_start, y_start, x_end, y_end]) + return click_map, bounding_boxes + + def get_patches_and_signals(self, img, click_map, bounding_boxes, cx, cy, m, n, bb=128): + + total = len(bounding_boxes) + img = np.array([img]) + click_map = np.array([click_map]) + click_map = click_map[:, np.newaxis, ...] + + patches = np.ndarray((total, 3, bb, bb), dtype=np.uint8) + nuc_points = np.ndarray((total, 1, bb, bb), dtype=np.uint8) + other_points = np.ndarray((total, 1, bb, bb), dtype=np.uint8) + + x_del_indices = {i for i in range(len(cx)) if cx[i] >= n or cx[i] < 0} + y_del_indices = {i for i in range(len(cy)) if cy[i] >= m or cy[i] < 0} + del_indices = list(x_del_indices.union(y_del_indices)) + cx = np.delete(cx, del_indices) + cy = np.delete(cy, del_indices) + + for i in range(len(bounding_boxes)): + bounding_box = bounding_boxes[i] + x_start = bounding_box[0] + y_start = bounding_box[1] + x_end = bounding_box[2] + y_end = bounding_box[3] + + patches[i] = img[0, :, y_start : y_end + 1, x_start : x_end + 1] + + this_click_map = np.zeros((1, 1, m, n), dtype=np.uint8) + this_click_map[0, 0, cy[i], cx[i]] = 1 + + others_click_map = np.uint8((click_map - this_click_map) > 0) + + nuc_points[i] = this_click_map[0, :, y_start : y_end + 1, x_start : x_end + 1] + other_points[i] = others_click_map[0, :, y_start : y_end + 1, x_start : x_end + 1] + + return patches, nuc_points, other_points + + +class PostFilterLabeld(MapTransform): + """ + Performs Filtering of Labels on the predicted probability map + + Args: + thresh: probability threshold for classifying a pixel as a mask + min_size: min_size objects that will be removed from the image, refer skimage remove_small_objects + min_hole: min_hole that will be removed from the image, refer skimage remove_small_holes + do_reconstruction: Boolean Flag, Perform a morphological reconstruction of an image, refer skimage + allow_missing_keys: don't raise exception if key is missing. + """ + + def __init__( + self, + keys: KeysCollection, + nuc_points: str = NuclickKeys.NUC_POINTS.value, + bounding_boxes: str = NuclickKeys.BOUNDING_BOXES.value, + img_height: str = NuclickKeys.IMG_HEIGHT.value, + img_width: str = NuclickKeys.IMG_WIDTH.value, + thresh: float = 0.33, + min_size: int = 10, + min_hole: int = 30, + do_reconstruction: bool = False, + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) + self.nuc_points = nuc_points + self.bounding_boxes = bounding_boxes + self.img_height = img_height + self.img_width = img_width + + self.thresh = thresh + self.min_size = min_size + self.min_hole = min_hole + self.do_reconstruction = do_reconstruction + + def __call__(self, data): + d = dict(data) + + nuc_points = d[self.nuc_points] + bounding_boxes = d[self.bounding_boxes] + img_height = d[self.img_height] + img_width = d[self.img_width] + + for key in self.keys: + label = d[key].astype(np.uint8) + masks = self.post_processing( + label, + thresh=self.thresh, + min_size=self.min_size, + min_hole=self.min_hole, + do_reconstruction=self.do_reconstruction, + nuc_points=nuc_points, + ) + + d[key] = self.gen_instance_map(masks, bounding_boxes, img_height, img_width).astype(np.uint8) + return d + + def post_processing(self, preds, thresh=0.33, min_size=10, min_hole=30, do_reconstruction=False, nuc_points=None): + masks = preds > thresh + masks = morphology.remove_small_objects(masks, min_size=min_size) + masks = morphology.remove_small_holes(masks, area_threshold=min_hole) + if do_reconstruction: + for i in range(len(masks)): + this_mask = masks[i] + this_marker = nuc_points[i, 0, :, :] > 0 + + try: + this_mask = morphology.reconstruction(this_marker, this_mask, footprint=morphology.disk(1)) + masks[i] = np.array([this_mask]) + except BaseException: + print("Nuclei reconstruction error #" + str(i)) + return masks + + def gen_instance_map(self, masks, bounding_boxes, m, n, flatten=True): + instance_map = np.zeros((m, n), dtype=np.uint16) + for i in range(len(masks)): + this_bb = bounding_boxes[i] + this_mask_pos = np.argwhere(masks[i] > 0) + this_mask_pos[:, 0] = this_mask_pos[:, 0] + this_bb[1] + this_mask_pos[:, 1] = this_mask_pos[:, 1] + this_bb[0] + instance_map[this_mask_pos[:, 0], this_mask_pos[:, 1]] = 1 if flatten else i + 1 + return instance_map diff --git a/monai/apps/pathology/handlers/prob_map_producer.py b/monai/apps/pathology/handlers/prob_map_producer.py index 0615b44b8f..62507dc0cb 100644 --- a/monai/apps/pathology/handlers/prob_map_producer.py +++ b/monai/apps/pathology/handlers/prob_map_producer.py @@ -16,7 +16,7 @@ import numpy as np from monai.config import DtypeLike, IgniteInfo -from monai.utils import min_version, optional_import +from monai.utils import deprecated, min_version, optional_import Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") if TYPE_CHECKING: @@ -25,6 +25,10 @@ Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") +@deprecated( + since="0.8", + msg_suffix="use `monai.handler.ProbMapProducer` (with `monai.data.wsi_dataset.SlidingPatchWSIDataset`) instead.", +) class ProbMapProducer: """ Event handler triggered on completing every iteration to save the probability map diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py index a44dce1e3f..ce22b86d49 100644 --- a/monai/apps/pathology/transforms/spatial/array.py +++ b/monai/apps/pathology/transforms/spatial/array.py @@ -17,12 +17,13 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.transforms.transform import Randomizable, Transform -from monai.utils import convert_data_type, convert_to_dst_type +from monai.utils import convert_data_type, convert_to_dst_type, deprecated from monai.utils.enums import TransformBackends __all__ = ["SplitOnGrid", "TileOnGrid"] +@deprecated(since="0.8", msg_suffix="use `monai.transforms.GridSplit` instead.") class SplitOnGrid(Transform): """ Split the image into patches based on the provided grid shape. @@ -107,6 +108,7 @@ def get_params(self, image_size): return patch_size, steps +@deprecated(since="0.8", msg_suffix="use `monai.transforms.GridPatch` or `monai.transforms.RandGridPatch` instead.") class TileOnGrid(Randomizable, Transform): """ Tile the 2D image into patches on a grid and maintain a subset of it. @@ -209,7 +211,7 @@ def __call__(self, image: NdarrayOrTensor) -> NdarrayOrTensor: constant_values=self.background_val, ) - # extact tiles + # extract tiles x_step, y_step = self.step, self.step h_tile, w_tile = self.tile_size, self.tile_size c_image, h_image, w_image = img_np.shape diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py index d5c34a0840..022d82a053 100644 --- a/monai/apps/pathology/transforms/spatial/dictionary.py +++ b/monai/apps/pathology/transforms/spatial/dictionary.py @@ -15,12 +15,14 @@ from monai.config import KeysCollection from monai.config.type_definitions import NdarrayOrTensor from monai.transforms.transform import MapTransform, Randomizable +from monai.utils import deprecated from .array import SplitOnGrid, TileOnGrid __all__ = ["SplitOnGridd", "SplitOnGridD", "SplitOnGridDict", "TileOnGridd", "TileOnGridD", "TileOnGridDict"] +@deprecated(since="0.8", msg_suffix="use `monai.transforms.GridSplitd` instead.") class SplitOnGridd(MapTransform): """ Split the image into patches based on the provided grid shape. @@ -55,6 +57,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d +@deprecated(since="0.8", msg_suffix="use `monai.transforms.GridPatchd` or `monai.transforms.RandGridPatchd` instead.") class TileOnGridd(Randomizable, MapTransform): """ Tile the 2D image into patches on a grid and maintain a subset of it. diff --git a/monai/apps/utils.py b/monai/apps/utils.py index 209dc796cf..cbfdcd7423 100644 --- a/monai/apps/utils.py +++ b/monai/apps/utils.py @@ -27,7 +27,7 @@ from monai.config.type_definitions import PathLike from monai.utils import look_up_option, min_version, optional_import -gdown, has_gdown = optional_import("gdown", "3.6") +gdown, has_gdown = optional_import("gdown", "4.4") if TYPE_CHECKING: from tqdm import tqdm @@ -147,7 +147,12 @@ def check_hash(filepath: PathLike, val: Optional[str] = None, hash_type: str = " def download_url( - url: str, filepath: PathLike = "", hash_val: Optional[str] = None, hash_type: str = "md5", progress: bool = True + url: str, + filepath: PathLike = "", + hash_val: Optional[str] = None, + hash_type: str = "md5", + progress: bool = True, + **gdown_kwargs, ) -> None: """ Download file from specified URL link, support process bar and hash check. @@ -160,6 +165,10 @@ def download_url( if None, skip hash validation. hash_type: 'md5' or 'sha1', defaults to 'md5'. 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. + details of the args of it: + https://github.com/wkentaro/gdown/blob/main/gdown/download.py Raises: RuntimeError: When the hash validation of the ``filepath`` existing file fails. @@ -189,7 +198,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.") - gdown.download(url, f"{tmp_name}", quiet=not progress) + gdown.download(url, f"{tmp_name}", quiet=not progress, **gdown_kwargs) else: _download_with_progress(url, tmp_name, progress=progress) if not tmp_name.exists(): diff --git a/monai/bundle/__init__.py b/monai/bundle/__init__.py index d6a452b5a4..2a658a3c51 100644 --- a/monai/bundle/__init__.py +++ b/monai/bundle/__init__.py @@ -12,5 +12,5 @@ from .config_item import ComponentLocator, ConfigComponent, ConfigExpression, ConfigItem, Instantiable from .config_parser import ConfigParser from .reference_resolver import ReferenceResolver -from .scripts import ckpt_export, run, verify_metadata, verify_net_in_out -from .utils import EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY +from .scripts import ckpt_export, download, init_bundle, load, run, verify_metadata, verify_net_in_out +from .utils import EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY, load_bundle_config diff --git a/monai/bundle/__main__.py b/monai/bundle/__main__.py index d77b396e79..ace3701d19 100644 --- a/monai/bundle/__main__.py +++ b/monai/bundle/__main__.py @@ -10,7 +10,7 @@ # limitations under the License. -from monai.bundle.scripts import ckpt_export, run, verify_metadata, verify_net_in_out +from monai.bundle.scripts import ckpt_export, download, init_bundle, run, verify_metadata, verify_net_in_out if __name__ == "__main__": from monai.utils import optional_import diff --git a/monai/bundle/config_item.py b/monai/bundle/config_item.py index 3300fe91ff..b410e5f641 100644 --- a/monai/bundle/config_item.py +++ b/monai/bundle/config_item.py @@ -281,7 +281,10 @@ def instantiate(self, **kwargs) -> object: # type: ignore modname = self.resolve_module_name() args = self.resolve_args() args.update(kwargs) - return instantiate(modname, **args) + try: + return instantiate(modname, **args) + except Exception as e: + raise RuntimeError(f"Failed to instantiate {self}.") from e class ConfigExpression(ConfigItem): @@ -312,7 +315,7 @@ class ConfigExpression(ConfigItem): """ prefix = EXPR_KEY - run_eval = False if os.environ.get("MONAI_EVAL_EXPR", "1") == "0" else True + run_eval = not os.environ.get("MONAI_EVAL_EXPR", "1") == "0" def __init__(self, config: Any, id: str = "", globals: Optional[Dict] = None) -> None: super().__init__(config=config, id=id) @@ -337,12 +340,13 @@ def _parse_import_string(self, import_string: str): return self.globals[asname] return None - def evaluate(self, locals: Optional[Dict] = None): + def evaluate(self, globals: Optional[Dict] = None, locals: Optional[Dict] = None): """ Execute the current config content and return the result if it is expression, based on Python `eval()`. For more details: https://docs.python.org/3/library/functions.html#eval. Args: + globals: besides ``self.globals``, other global symbols used in the expression at runtime. locals: besides ``globals``, may also have some local symbols used in the expression at runtime. """ @@ -354,7 +358,13 @@ def evaluate(self, locals: Optional[Dict] = None): return optional_module if not self.run_eval: return f"{value[len(self.prefix) :]}" - return eval(value[len(self.prefix) :], self.globals, locals) + globals_ = dict(self.globals) + if globals is not None: + for k, v in globals.items(): + if k in globals_: + warnings.warn(f"the new global variable `{k}` conflicts with `self.globals`, override it.") + globals_[k] = v + return eval(value[len(self.prefix) :], globals_, locals) @classmethod def is_expression(cls, config: Union[Dict, List, str]) -> bool: diff --git a/monai/bundle/config_parser.py b/monai/bundle/config_parser.py index 4f919a383a..791759fd65 100644 --- a/monai/bundle/config_parser.py +++ b/monai/bundle/config_parser.py @@ -183,6 +183,19 @@ def set(self, config: Any, id: str = ""): """ self[id] = config + def __contains__(self, id: Union[str, int]) -> bool: + """ + Returns True if `id` is stored in this configuration. + + Args: + id: id to specify the expected position. See also :py:meth:`__getitem__`. + """ + try: + _ = self[id] + return True + except KeyError: + return False + def parse(self, reset: bool = True): """ Recursively resolve `self.config` to replace the macro tokens with target content. @@ -211,16 +224,17 @@ def get_parsed_content(self, id: str = "", **kwargs): Use digits indexing from "0" for list or other strings for dict. For example: ``"xform#5"``, ``"net#channels"``. ``""`` indicates the entire ``self.config``. kwargs: additional keyword arguments to be passed to ``_resolve_one_item``. - Currently support ``lazy`` (whether to retain the current config cache, default to `False`), + Currently support ``lazy`` (whether to retain the current config cache, default to `True`), ``instantiate`` (whether to instantiate the `ConfigComponent`, default to `True`) and - ``eval_expr`` (whether to evaluate the `ConfigExpression`, default to `True`). + ``eval_expr`` (whether to evaluate the `ConfigExpression`, default to `True`), ``default`` + (the default config item if the `id` is not in the config content). """ if not self.ref_resolver.is_resolved(): # not parsed the config source yet, parse it self.parse(reset=True) - elif not kwargs.get("lazy", False): - self.parse(reset=not kwargs.get("lazy", False)) + elif not kwargs.get("lazy", True): + self.parse(reset=not kwargs.get("lazy", True)) return self.ref_resolver.get_resolved_content(id=id, **kwargs) def read_meta(self, f: Union[PathLike, Sequence[PathLike], Dict], **kwargs): @@ -230,7 +244,7 @@ def read_meta(self, f: Union[PathLike, Sequence[PathLike], Dict], **kwargs): Args: f: filepath of the metadata file, the content must be a dictionary, - if providing a list of files, wil merge the content of them. + if providing a list of files, will merge the content of them. if providing a dictionary directly, use it as metadata. kwargs: other arguments for ``json.load`` or ``yaml.safe_load``, depends on the file format. @@ -338,9 +352,12 @@ def load_config_file(cls, filepath: PathLike, **kwargs): raise ValueError(f"only support JSON or YAML config file so far, got name {_filepath}.") @classmethod - def load_config_files(cls, files: Union[PathLike, Sequence[PathLike], dict], **kwargs) -> dict: + def load_config_files(cls, files: Union[PathLike, Sequence[PathLike], dict], **kwargs) -> Dict: """ Load config files into a single config dict. + The latter config file in the list will override or add the former config file. + ``"#"`` in the config keys are interpreted as special characters to go one level + further into the nested structures. Args: files: path of target files to load, supported postfixes: `.json`, `.yml`, `.yaml`. @@ -348,10 +365,11 @@ def load_config_files(cls, files: Union[PathLike, Sequence[PathLike], dict], **k """ if isinstance(files, dict): # already a config dict return files - content = {} + parser = ConfigParser(config={}) for i in ensure_tuple(files): - content.update(cls.load_config_file(i, **kwargs)) - return content + for k, v in (cls.load_config_file(i, **kwargs)).items(): + parser[k] = v + return parser.get() # type: ignore @classmethod def export_config_file(cls, config: Dict, filepath: PathLike, fmt="json", **kwargs): diff --git a/monai/bundle/reference_resolver.py b/monai/bundle/reference_resolver.py index f9f73c9c71..7ed5ae3522 100644 --- a/monai/bundle/reference_resolver.py +++ b/monai/bundle/reference_resolver.py @@ -9,7 +9,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import re +import warnings from typing import Any, Dict, Optional, Sequence, Set from monai.bundle.config_item import ConfigComponent, ConfigExpression, ConfigItem @@ -50,6 +52,8 @@ class ReferenceResolver: ref = ID_REF_KEY # reference prefix # match a reference string, e.g. "@id#key", "@id#key#0", "@_target_#key" id_matcher = re.compile(rf"{ref}(?:\w*)(?:{sep}\w*)*") + # if `allow_missing_reference` and can't find a reference ID, will just raise a warning and don't update the config + allow_missing_reference = not os.environ.get("MONAI_ALLOW_MISSING_REFERENCE", "0") == "0" def __init__(self, items: Optional[Sequence[ConfigItem]] = None): # save the items in a dictionary with the `ConfigItem.id` as key @@ -107,14 +111,16 @@ def _resolve_one_item(self, id: str, waiting_list: Optional[Set[str]] = None, ** waiting_list: set of ids pending to be resolved. It's used to detect circular references such as: `{"name": "A", "dep": "@B"}` and `{"name": "B", "dep": "@A"}`. - kwargs: keyword arguments to pass to ``_resolve_one_item()``. - Currently support ``instantiate`` and ``eval_expr``. Both are defaulting to True. + kwargs: keyword arguments to pass to ``_resolve_one_item()``. + Currently support ``instantiate``, ``eval_expr`` and ``default``. + `instantiate` and `eval_expr` are defaulting to True, `default` is the target config item + if the `id` is not in the config content, must be a `ConfigItem` object. """ if id in self.resolved_content: return self.resolved_content[id] try: - item = look_up_option(id, self.items, print_all_options=False) + item = look_up_option(id, self.items, print_all_options=False, default=kwargs.get("default", "no_default")) except ValueError as err: raise KeyError(f"id='{id}' is not found in the config resolver.") from err item_config = item.get_config() @@ -130,7 +136,7 @@ def _resolve_one_item(self, id: str, waiting_list: Optional[Set[str]] = None, ** and v.is_import_statement(v.get_config()) ): self.resolved_content[t] = v.evaluate() if kwargs.get("eval_expr", True) else v - for d in self.find_refs_in_config(config=item_config, id=id): + for d in self.find_refs_in_config(config=item_config, id=id).keys(): # if current item has reference already in the waiting list, that's circular references if d in waiting_list: raise ValueError(f"detected circular references '{d}' for id='{id}' in the config content.") @@ -140,7 +146,12 @@ def _resolve_one_item(self, id: str, waiting_list: Optional[Set[str]] = None, ** try: look_up_option(d, self.items, print_all_options=False) except ValueError as err: - raise ValueError(f"the referring item `@{d}` is not defined in the config content.") from err + msg = f"the referring item `@{d}` is not defined in the config content." + if self.allow_missing_reference: + warnings.warn(msg) + continue + else: + raise ValueError(msg) from err # recursively resolve the reference first self._resolve_one_item(id=d, waiting_list=waiting_list, **kwargs) waiting_list.discard(d) @@ -154,7 +165,7 @@ def _resolve_one_item(self, id: str, waiting_list: Optional[Set[str]] = None, ** elif isinstance(item, ConfigExpression): run_eval = kwargs.get("eval_expr", True) self.resolved_content[id] = ( - item.evaluate(locals={f"{self._vars}": self.resolved_content}) if run_eval else item + item.evaluate(globals={f"{self._vars}": self.resolved_content}) if run_eval else item ) else: self.resolved_content[id] = new_config @@ -166,14 +177,16 @@ def get_resolved_content(self, id: str, **kwargs): Args: id: id name of the expected item. - kwargs: additional keyword arguments to be passed to ``_resolve_one_item``. - Currently support ``instantiate`` and ``eval_expr``. Both are defaulting to True. + kwargs: keyword arguments to pass to ``_resolve_one_item()``. + Currently support ``instantiate``, ``eval_expr`` and ``default``. + `instantiate` and `eval_expr` are defaulting to True, `default` is the target config item + if the `id` is not in the config content, must be a `ConfigItem` object. """ return self._resolve_one_item(id=id, **kwargs) @classmethod - def match_refs_pattern(cls, value: str) -> Set[str]: + def match_refs_pattern(cls, value: str) -> Dict[str, int]: """ Match regular expression for the input string to find the references. The reference string starts with ``"@"``, like: ``"@XXX#YYY#ZZZ"``. @@ -182,14 +195,15 @@ def match_refs_pattern(cls, value: str) -> Set[str]: value: input value to match regular expression. """ - refs: Set[str] = set() + refs: Dict[str, int] = {} # regular expression pattern to match "@XXX" or "@XXX#YYY" result = cls.id_matcher.findall(value) value_is_expr = ConfigExpression.is_expression(value) for item in result: if value_is_expr or value == item: # only check when string starts with "$" or the whole content is "@XXX" - refs.add(item[len(cls.ref) :]) + id = item[len(cls.ref) :] + refs[id] = refs.get(id, 0) + 1 return refs @classmethod @@ -210,9 +224,15 @@ def update_refs_pattern(cls, value: str, refs: Dict) -> str: for item in result: ref_id = item[len(cls.ref) :] # remove the ref prefix "@" if ref_id not in refs: - raise KeyError(f"can not find expected ID '{ref_id}' in the references.") + msg = f"can not find expected ID '{ref_id}' in the references." + if cls.allow_missing_reference: + warnings.warn(msg) + continue + else: + raise KeyError(msg) if value_is_expr: - # replace with local code, will be used in the `evaluate` logic with `locals={"refs": ...}` + # replace with local code, `{"__local_refs": self.resolved_content}` will be added to + # the `globals` argument of python `eval` in the `evaluate` value = value.replace(item, f"{cls._vars}['{ref_id}']") elif value == item: # the whole content is "@XXX", it will avoid the case that regular string contains "@" @@ -220,7 +240,7 @@ def update_refs_pattern(cls, value: str, refs: Dict) -> str: return value @classmethod - def find_refs_in_config(cls, config, id: str, refs: Optional[Set[str]] = None) -> Set[str]: + def find_refs_in_config(cls, config, id: str, refs: Optional[Dict[str, int]] = None) -> Dict[str, int]: """ Recursively search all the content of input config item to get the ids of references. References mean: the IDs of other config items (``"@XXX"`` in this config item), or the @@ -230,18 +250,19 @@ def find_refs_in_config(cls, config, id: str, refs: Optional[Set[str]] = None) - Args: config: input config content to search. id: ID name for the input config item. - refs: list of the ID name of found references, default to `None`. + refs: dict of the ID name and count of found references, default to `None`. """ - refs_: Set[str] = refs or set() + refs_: Dict[str, int] = refs or {} if isinstance(config, str): - return refs_.union(cls.match_refs_pattern(value=config)) + for id, count in cls.match_refs_pattern(value=config).items(): + refs_[id] = refs_.get(id, 0) + count if not isinstance(config, (list, dict)): return refs_ 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}" - if ConfigComponent.is_instantiable(v) or ConfigExpression.is_expression(v): - refs_.add(sub_id) + if ConfigComponent.is_instantiable(v) or ConfigExpression.is_expression(v) and sub_id not in refs_: + refs_[sub_id] = 1 refs_ = cls.find_refs_in_config(v, sub_id, refs_) return refs_ diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 64172c4541..ff8144397f 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -11,24 +11,33 @@ import ast import json +import os import pprint import re +import warnings from logging.config import fileConfig -from typing import Dict, Optional, Sequence, Tuple, Union +from pathlib import Path +from shutil import copyfile +from textwrap import dedent +from typing import Dict, Mapping, Optional, Sequence, Tuple, Union import torch from torch.cuda import is_available -from monai.apps.utils import download_url, get_logger +from monai.apps.utils import _basename, download_url, extractall, get_logger +from monai.bundle.config_item import ConfigComponent from monai.bundle.config_parser import ConfigParser +from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA from monai.config import IgniteInfo, PathLike -from monai.data import save_net_with_metadata -from monai.networks import convert_to_torchscript, copy_model_state +from monai.data import load_net_with_metadata, save_net_with_metadata +from monai.networks import convert_to_torchscript, copy_model_state, get_state_dict, save_state from monai.utils import check_parent_dir, get_equivalent_dtype, min_version, optional_import +from monai.utils.misc import ensure_tuple 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_get, has_requests = optional_import("requests", name="get") logger = get_logger(module_name=__name__) @@ -44,7 +53,7 @@ def _update_args(args: Optional[Union[str, Dict]] = None, ignore_none: bool = Tr kwargs: destination args to update. """ - args_: Dict = args if isinstance(args, dict) else {} # type: ignore + args_: Dict = args if isinstance(args, dict) else {} if isinstance(args, str): # args are defined in a structured file args_ = ConfigParser.load_config_file(args) @@ -116,8 +125,186 @@ def _get_fake_spatial_shape(shape: Sequence[Union[str, int]], p: int = 1, n: int return tuple(ret) +def _get_git_release_url(repo_owner: str, repo_name: str, tag_name: str, filename: str): + return f"https://github.com/{repo_owner}/{repo_name}/releases/download/{tag_name}/{filename}" + + +def _download_from_github(repo: str, download_path: Path, filename: str, progress: bool = True): + if len(repo.split("/")) != 3: + raise ValueError("if source is `github`, repo should be in the form of `repo_owner/repo_name/release_tag`.") + repo_owner, repo_name, tag_name = repo.split("/") + if ".zip" not in filename: + filename += ".zip" + url = _get_git_release_url(repo_owner, repo_name, tag_name=tag_name, filename=filename) + filepath = download_path / f"{filename}" + download_url(url=url, filepath=filepath, hash_val=None, progress=progress) + extractall(filepath=filepath, output_dir=download_path, has_base=True) + + +def _process_bundle_dir(bundle_dir: Optional[PathLike] = None): + if bundle_dir is None: + get_dir, has_home = optional_import("torch.hub", name="get_dir") + if has_home: + bundle_dir = Path(get_dir()) / "bundle" + else: + raise ValueError("bundle_dir=None, but no suitable default directory computed. Upgrade Pytorch to 1.6+ ?") + return Path(bundle_dir) + + +def download( + name: Optional[str] = None, + bundle_dir: Optional[PathLike] = None, + source: str = "github", + repo: str = "Project-MONAI/model-zoo/hosting_storage_v1", + url: Optional[str] = None, + progress: bool = True, + args_file: Optional[str] = None, +): + """ + download bundle from the specified source or url. The bundle should be a zip file and it + will be extracted after downloading. + This function refers to: + https://pytorch.org/docs/stable/_modules/torch/hub.html + + Typical usage examples: + + .. code-block:: bash + + # Execute this module as a CLI entry, and download bundle: + python -m monai.bundle download --name "bundle_name" --source "github" --repo "repo_owner/repo_name/release_tag" + + # Execute this module as a CLI entry, and download bundle via URL: + python -m monai.bundle download --name "bundle_name" --url + + # Set default args of `run` in a JSON / YAML file, help to record and simplify the command line. + # Other args still can override the default args at runtime. + # The content of the JSON / YAML file is a dictionary. For example: + # {"name": "spleen", "bundle_dir": "download", "source": ""} + # then do the following command for downloading: + python -m monai.bundle download --args_file "args.json" --source "github" + + Args: + name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`. + bundle_dir: target directory to store the downloaded data. + Default is `bundle` subfolder under `torch.hub.get_dir()`. + source: storage location name. This argument is used when `url` is `None`. + "github" is currently the only supported value. + repo: repo name. This argument is used when `url` is `None`. + If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag". + url: url to download the data. If not `None`, data will be downloaded directly + and `source` will not be checked. + If `name` is `None`, filename is determined by `monai.apps.utils._basename(url)`. + progress: whether to display a progress bar. + args_file: a JSON or YAML file to provide default values for all the args in this function. + so that the command line inputs can be simplified. + + """ + _args = _update_args( + args=args_file, name=name, bundle_dir=bundle_dir, source=source, repo=repo, url=url, progress=progress + ) + + _log_input_summary(tag="download", args=_args) + source_, repo_, progress_, name_, bundle_dir_, url_ = _pop_args( + _args, "source", "repo", "progress", name=None, bundle_dir=None, url=None + ) + + bundle_dir_ = _process_bundle_dir(bundle_dir_) + + if url_ is not None: + if name is not None: + filepath = bundle_dir_ / f"{name}.zip" + else: + filepath = bundle_dir_ / f"{_basename(url_)}" + download_url(url=url_, filepath=filepath, hash_val=None, progress=progress_) + extractall(filepath=filepath, output_dir=bundle_dir_, has_base=True) + elif source_ == "github": + if name_ is None: + raise ValueError(f"To download from source: Github, `name` must be provided, got {name_}.") + _download_from_github(repo=repo_, download_path=bundle_dir_, filename=name_, progress=progress_) + else: + raise NotImplementedError( + f"Currently only download from provided URL in `url` or Github is implemented, got source: {source_}." + ) + + +def load( + name: str, + model_file: Optional[str] = None, + load_ts_module: bool = False, + bundle_dir: Optional[PathLike] = None, + source: str = "github", + repo: str = "Project-MONAI/model-zoo/hosting_storage_v1", + progress: bool = True, + device: Optional[str] = None, + key_in_ckpt: Optional[str] = None, + config_files: Sequence[str] = (), + net_name: Optional[str] = None, + **net_kwargs, +): + """ + Load model weights or TorchScript module of a bundle. + + Args: + name: bundle name. + model_file: the relative path of the model weights or TorchScript module within bundle. + If `None`, "models/model.pt" or "models/model.ts" will be used. + load_ts_module: a flag to specify if loading the TorchScript module. + bundle_dir: directory the weights/TorchScript module will be loaded from. + Default is `bundle` subfolder under `torch.hub.get_dir()`. + source: storage location name. This argument is used when `model_file` is not existing locally and need to be + downloaded first. "github" is currently the only supported value. + repo: repo name. This argument is used when `model_file` is not existing locally and need to be + downloaded first. If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag". + progress: whether to display a progress bar when downloading. + device: target device of returned weights or module, if `None`, prefer to "cuda" if existing. + key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model + weights. if not nested checkpoint, no need to set. + config_files: extra filenames would be loaded. The argument only works when loading a TorchScript module, + see `_extra_files` in `torch.jit.load` for more details. + net_name: if not `None`, a corresponding network will be instantiated and load the achieved weights. + This argument only works when loading weights. + net_kwargs: other arguments that are used to instantiate the network class defined by `net_name`. + + Returns: + 1. If `load_ts_module` is `False` and `net_name` is `None`, return model weights. + 2. If `load_ts_module` is `False` and `net_name` is not `None`, + return an instantiated network that loaded the weights. + 3. If `load_ts_module` is `True`, return a triple that include a TorchScript module, + the corresponding metadata dict, and extra files dict. + please check `monai.data.load_net_with_metadata` for more details. + + """ + bundle_dir_ = _process_bundle_dir(bundle_dir) + + if model_file is None: + model_file = os.path.join("models", "model.ts" if load_ts_module is True else "model.pt") + full_path = os.path.join(bundle_dir_, name, model_file) + if not os.path.exists(full_path): + download(name=name, bundle_dir=bundle_dir_, source=source, repo=repo, progress=progress) + + if device is None: + device = "cuda:0" if is_available() else "cpu" + # loading with `torch.jit.load` + if load_ts_module is True: + return load_net_with_metadata(full_path, map_location=torch.device(device), more_extra_files=config_files) + # loading with `torch.load` + model_dict = torch.load(full_path, map_location=torch.device(device)) + if not isinstance(model_dict, Mapping): + warnings.warn(f"the state dictionary from {full_path} should be a dictionary but got {type(model_dict)}.") + model_dict = get_state_dict(model_dict) + + if net_name is None: + return model_dict + net_kwargs["_target_"] = net_name + configer = ConfigComponent(config=net_kwargs) + model = configer.instantiate() + model.to(device) # type: ignore + copy_model_state(dst=model, src=model_dict if key_in_ckpt is None else model_dict[key_in_ckpt]) # type: ignore + return model + + def run( - runner_id: Optional[str] = None, + runner_id: Optional[Union[str, Sequence[str]]] = None, meta_file: Optional[Union[str, Sequence[str]]] = None, config_file: Optional[Union[str, Sequence[str]]] = None, logging_file: Optional[str] = None, @@ -132,23 +319,23 @@ def run( .. code-block:: bash # Execute this module as a CLI entry: - python -m monai.bundle run trainer --meta_file --config_file + python -m monai.bundle run training --meta_file --config_file # Override config values at runtime by specifying the component id and its new value: - python -m monai.bundle run trainer --net#input_chns 1 ... + python -m monai.bundle run training --net#input_chns 1 ... # Override config values with another config file `/path/to/another.json`: - python -m monai.bundle run evaluator --net %/path/to/another.json ... + python -m monai.bundle run evaluating --net %/path/to/another.json ... # Override config values with part content of another config file: - python -m monai.bundle run trainer --net %/data/other.json#net_arg ... + python -m monai.bundle run training --net %/data/other.json#net_arg ... # Set default args of `run` in a JSON / YAML file, help to record and simplify the command line. # Other args still can override the default args at runtime: python -m monai.bundle run --args_file "/workspace/data/args.json" --config_file Args: - runner_id: ID name of the runner component or workflow, it must have a `run` method. Defaults to ``""``. + runner_id: ID name of the expected config expression to run, can also be a list of IDs to run in order. meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged. config_file: filepath of the config file, if `None`, must be provided in `args_file`. if it is a list of file paths, the content of them will be merged. @@ -188,12 +375,8 @@ def run( for k, v in _args.items(): parser[k] = v - workflow = parser.get_parsed_content(id=runner_id_) - if not hasattr(workflow, "run"): - raise ValueError( - f"The parsed workflow {type(workflow)} (id={runner_id_}) does not have a `run` method.\n{run.__doc__}" - ) - return workflow.run() + # resolve and execute the specified runner expressions in the config, return the results + return [parser.get_parsed_content(i, lazy=True, eval_expr=True, instantiate=True) for i in ensure_tuple(runner_id_)] def verify_metadata( @@ -249,7 +432,7 @@ def verify_metadata( try: # the rest key-values in the _args are for `validate` API validate(instance=metadata, schema=schema, **_args) - except ValidationError as e: + except ValidationError as e: # pylint: disable=E0712 # as the error message is very long, only extract the key information logger.info(re.compile(r".*Failed validating", re.S).findall(str(e))[0] + f" against schema `{url}`.") return @@ -336,7 +519,7 @@ def verify_net_in_out( net.eval() with torch.no_grad(): - spatial_shape = _get_fake_spatial_shape(input_spatial_shape, p=p_, n=n_, any=any_) # type: ignore + spatial_shape = _get_fake_spatial_shape(input_spatial_shape, p=p_, n=n_, any=any_) test_data = torch.rand(*(1, input_channels, *spatial_shape), dtype=input_dtype, device=device_) output = net(test_data) if output.shape[1] != output_channels: @@ -363,15 +546,17 @@ def ckpt_export( .. code-block:: bash - python -m monai.bundle export network --filepath --ckpt_file ... + python -m monai.bundle ckpt_export network --filepath --ckpt_file ... Args: net_id: ID name of the network component in the config, it must be `torch.nn.Module`. filepath: filepath to export, if filename has no extension it becomes `.ts`. ckpt_file: filepath of the model checkpoint to load. meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged. - config_file: filepath of the config file, if `None`, must be provided in `args_file`. - if it is a list of file paths, the content of them will be merged. + config_file: filepath of the config file to save in TorchScript model and extract network information, + the saved key in the TorchScript model is the config filename without extension, and the saved config + value is always serialized in JSON format no matter the original file format is JSON or YAML. + it can be a single file or a list of files. if `None`, must be provided in `args_file`. key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model weights. if not nested checkpoint, no need to set. args_file: a JSON or YAML file to provide default values for `meta_file`, `config_file`, @@ -390,7 +575,7 @@ def ckpt_export( key_in_ckpt=key_in_ckpt, **override, ) - _log_input_summary(tag="export", args=_args) + _log_input_summary(tag="ckpt_export", args=_args) filepath_, ckpt_file_, config_file_, net_id_, meta_file_, key_in_ckpt_ = _pop_args( _args, "filepath", "ckpt_file", "config_file", net_id="", meta_file=None, key_in_ckpt="" ) @@ -410,17 +595,110 @@ def ckpt_export( # here we use ignite Checkpoint to support nested weights and be compatible with MONAI CheckpointSaver Checkpoint.load_objects(to_load={key_in_ckpt_: net}, checkpoint=ckpt_file_) else: - copy_model_state(dst=net, src=ckpt_file_ if key_in_ckpt_ == "" else ckpt_file_[key_in_ckpt_]) + ckpt = torch.load(ckpt_file_) + copy_model_state(dst=net, src=ckpt if key_in_ckpt_ == "" else ckpt[key_in_ckpt_]) - # convert to TorchScript model and save with meta data, config content + # convert to TorchScript model and save with metadata, config content net = convert_to_torchscript(model=net) + extra_files: Dict = {} + for i in ensure_tuple(config_file_): + # split the filename and directory + filename = os.path.basename(i) + # remove extension + filename, _ = os.path.splitext(filename) + # because all files are stored as JSON their name parts without extension must be unique + if filename in extra_files: + raise ValueError(f"Filename part '{filename}' is given multiple times in config file list.") + # the file may be JSON or YAML but will get loaded and dumped out again as JSON + extra_files[filename] = json.dumps(ConfigParser.load_config_file(i)).encode() + + # add .json extension to all extra files which are always encoded as JSON + extra_files = {k + ".json": v for k, v in extra_files.items()} + save_net_with_metadata( jit_obj=net, filename_prefix_or_stream=filepath_, include_config_vals=False, append_timestamp=False, meta_values=parser.get().pop("_meta_", None), - more_extra_files={"config": json.dumps(parser.get()).encode()}, + more_extra_files=extra_files, ) logger.info(f"exported to TorchScript file: {filepath_}.") + + +def init_bundle( + bundle_dir: PathLike, + ckpt_file: Optional[PathLike] = None, + network: Optional[torch.nn.Module] = None, + metadata_str: Union[Dict, str] = DEFAULT_METADATA, + inference_str: Union[Dict, str] = DEFAULT_INFERENCE, +): + """ + Initialise a new bundle directory with some default configuration files and optionally network weights. + + Typical usage example: + + .. code-block:: bash + + python -m monai.bundle init_bundle /path/to/bundle_dir network_ckpt.pt + + Args: + bundle_dir: directory name to create, must not exist but parent direct must exist + ckpt_file: optional checkpoint file to copy into bundle + network: if given instead of ckpt_file this network's weights will be stored in bundle + """ + + bundle_dir = Path(bundle_dir).absolute() + + if bundle_dir.exists(): + raise ValueError(f"Specified bundle directory '{str(bundle_dir)}' already exists") + + if not bundle_dir.parent.is_dir(): + raise ValueError(f"Parent directory of specified bundle directory '{str(bundle_dir)}' does not exist") + + configs_dir = bundle_dir / "configs" + models_dir = bundle_dir / "models" + docs_dir = bundle_dir / "docs" + + bundle_dir.mkdir() + configs_dir.mkdir() + models_dir.mkdir() + docs_dir.mkdir() + + if isinstance(metadata_str, dict): + metadata_str = json.dumps(metadata_str, indent=4) + + if isinstance(inference_str, dict): + inference_str = json.dumps(inference_str, indent=4) + + with open(str(configs_dir / "metadata.json"), "w") as o: + o.write(metadata_str) + + with open(str(configs_dir / "inference.json"), "w") as o: + o.write(inference_str) + + with open(str(docs_dir / "README.md"), "w") as o: + readme = """ + # Your Model Name + + Describe your model here and how to run it, for example using `inference.json`: + + ``` + python -m monai.bundle run evaluating \ + --meta_file /path/to/bundle/configs/metadata.json \ + --config_file /path/to/bundle/configs/inference.json \ + --dataset_dir ./input \ + --bundle_root /path/to/bundle + ``` + """ + + o.write(dedent(readme)) + + with open(str(docs_dir / "license.txt"), "w") as o: + o.write("Select a license and place its terms here\n") + + if ckpt_file is not None: + copyfile(str(ckpt_file), str(models_dir / "model.pt")) + elif network is not None: + save_state(network, str(models_dir / "model.pt")) diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index ba5c2729e7..e3eb362941 100644 --- a/monai/bundle/utils.py +++ b/monai/bundle/utils.py @@ -9,6 +9,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import os +import zipfile +from typing import Any + +from monai.config.deviceconfig import get_config_values +from monai.utils import optional_import + +yaml, _ = optional_import("yaml") + __all__ = ["ID_REF_KEY", "ID_SEP_KEY", "EXPR_KEY", "MACRO_KEY"] @@ -16,3 +26,147 @@ ID_SEP_KEY = "#" # separator for the ID of a ConfigItem EXPR_KEY = "$" # start of a ConfigExpression MACRO_KEY = "%" # start of a macro of a config + + +_conf_values = get_config_values() + +DEFAULT_METADATA = { + "version": "0.0.1", + "changelog": {"0.0.1": "Initial version"}, + "monai_version": _conf_values["MONAI"], + "pytorch_version": _conf_values["Pytorch"], + "numpy_version": _conf_values["Numpy"], + "optional_packages_version": {}, + "task": "Describe what the network predicts", + "description": "A longer description of what the network does, use context, inputs, outputs, etc.", + "authors": "Your Name Here", + "copyright": "Copyright (c) Your Name Here", + "network_data_format": {"inputs": {}, "outputs": {}}, +} + +DEFAULT_INFERENCE = { + "imports": ["$import glob"], + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "ckpt_path": "$@bundle_root + '/models/model.pt'", + "dataset_dir": "/workspace/data", + "datalist": "$list(sorted(glob.glob(@dataset_dir + '/*.jpeg')))", + "network_def": {"_target_": "???", "spatial_dims": 2}, + "network": "$@network_def.to(@device)", + "preprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "LoadImaged", "keys": "image"}, + {"_target_": "AddChanneld", "keys": "image"}, + {"_target_": "ScaleIntensityd", "keys": "image"}, + {"_target_": "EnsureTyped", "keys": "image", "device": "@device"}, + ], + }, + "dataset": {"_target_": "Dataset", "data": "$[{'image': i} for i in @datalist]", "transform": "@preprocessing"}, + "dataloader": { + "_target_": "DataLoader", + "dataset": "@dataset", + "batch_size": 1, + "shuffle": False, + "num_workers": 0, + }, + "inferer": {"_target_": "SimpleInferer"}, + "postprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "Activationsd", "keys": "pred", "softmax": True}, + {"_target_": "AsDiscreted", "keys": "pred", "argmax": True}, + ], + }, + "handlers": [ + { + "_target_": "CheckpointLoader", + "_disabled_": "$not os.path.exists(@ckpt_path)", + "load_path": "@ckpt_path", + "load_dict": {"model": "@network"}, + } + ], + "evaluator": { + "_target_": "SupervisedEvaluator", + "device": "@device", + "val_data_loader": "@dataloader", + "network": "@network", + "inferer": "@inferer", + "postprocessing": "@postprocessing", + "val_handlers": "@handlers", + }, + "evaluating": ["$@evaluator.run()"], +} + + +def load_bundle_config(bundle_path: str, *config_names, **load_kw_args) -> Any: + """ + Load the metadata and nominated configuration files from a MONAI bundle without loading the network itself. + + This function will load the information from the bundle, which can be a directory or a zip file containing a + directory or a Torchscript bundle, and return the parser object with the information. This saves having to load + the model if only the information is wanted, and can work on any sort of bundle format. + + Args: + bundle_path: path to the bundle directory or zip file + config_names: names of configuration files with extensions to load, should not be full paths but just name+ext + load_kw_args: keyword arguments to pass to the ConfigParser object when loading + + Returns: + ConfigParser object containing the parsed information + """ + + from monai.bundle.config_parser import ConfigParser # avoids circular import + + parser = ConfigParser() + + if not os.path.exists(bundle_path): + raise ValueError(f"Cannot find bundle file/directory '{bundle_path}'") + + # bundle is a directory, read files directly + if os.path.isdir(bundle_path): + conf_data = [] + parser.read_meta(f=os.path.join(bundle_path, "configs", "metadata.json"), **load_kw_args) + + for cname in config_names: + cpath = os.path.join(bundle_path, "configs", cname) + if not os.path.exists(cpath): + raise ValueError(f"Cannot find config file '{cpath}'") + + conf_data.append(cpath) + + parser.read_config(f=conf_data, **load_kw_args) + else: + # bundle is a zip file which is either a zipped directory or a Torchscript archive + + 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: + raise ValueError(f"Cannot find config file '{full_cname}'") + + ardata = archive.read(full_cname) + + if full_cname.lower().endswith("json"): + cdata = json.loads(ardata, **load_kw_args) + elif full_cname.lower().endswith(("yaml", "yml")): + cdata = yaml.safe_load(ardata, **load_kw_args) + + parser.read_config(f=cdata) + + return parser diff --git a/monai/config/__init__.py b/monai/config/__init__.py index bf1b66fe92..5f67ea6584 100644 --- a/monai/config/__init__.py +++ b/monai/config/__init__.py @@ -28,5 +28,6 @@ NdarrayOrTensor, NdarrayTensor, PathLike, + SequenceStr, TensorOrList, ) diff --git a/monai/config/deviceconfig.py b/monai/config/deviceconfig.py index fd7ca572e6..ad633a133d 100644 --- a/monai/config/deviceconfig.py +++ b/monai/config/deviceconfig.py @@ -75,6 +75,7 @@ def get_optional_config_values(): output["einops"] = get_package_version("einops") output["transformers"] = get_package_version("transformers") output["mlflow"] = get_package_version("mlflow") + output["pynrrd"] = get_package_version("nrrd") return output @@ -120,7 +121,8 @@ def get_system_info() -> OrderedDict: if output["System"] == "Windows": _dict_append(output, "Win32 version", platform.win32_ver) if hasattr(platform, "win32_edition"): - _dict_append(output, "Win32 edition", platform.win32_edition) # type:ignore[attr-defined] + _dict_append(output, "Win32 edition", platform.win32_edition) + elif output["System"] == "Darwin": _dict_append(output, "Mac version", lambda: platform.mac_ver()[0]) else: diff --git a/monai/config/type_definitions.py b/monai/config/type_definitions.py index 16919c2ec4..bb6f87e97a 100644 --- a/monai/config/type_definitions.py +++ b/monai/config/type_definitions.py @@ -38,6 +38,7 @@ "NdarrayOrTensor", "TensorOrList", "PathLike", + "SequenceStr", ] @@ -77,3 +78,7 @@ #: PathLike: The PathLike type is used for defining a file path. PathLike = Union[str, os.PathLike] + +#: SequenceStr +# string or a sequence of strings for `mode` types. +SequenceStr = Union[Sequence[str], str] diff --git a/monai/data/__init__.py b/monai/data/__init__.py index bed194d2f4..b7a160b3de 100644 --- a/monai/data/__init__.py +++ b/monai/data/__init__.py @@ -9,6 +9,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib + +from .box_utils import ( + box_area, + box_centers, + box_giou, + box_iou, + box_pair_giou, + boxes_center_distance, + centers_in_boxes, + convert_box_mode, + convert_box_to_standard_mode, +) from .csv_saver import CSVSaver from .dataloader import DataLoader from .dataset import ( @@ -32,9 +45,9 @@ load_decathlon_properties, ) from .folder_layout import FolderLayout -from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter +from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter, PatchIterd from .image_dataset import ImageDataset -from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader, WSIReader +from .image_reader import ImageReader, ITKReader, NibabelReader, NrrdReader, NumpyReader, PILReader, PydicomReader from .image_writer import ( SUPPORTED_WRITERS, ImageWriter, @@ -46,6 +59,8 @@ resolve_writer, ) from .iterable_dataset import CSVIterableDataset, IterableDataset, ShuffleBuffer +from .meta_obj import MetaObj, get_track_meta, set_track_meta +from .meta_tensor import MetaTensor from .nifti_saver import NiftiSaver from .nifti_writer import write_nifti from .png_saver import PNGSaver @@ -56,6 +71,7 @@ from .thread_buffer import ThreadBuffer, ThreadDataLoader from .torchscript_utils import load_net_with_metadata, save_net_with_metadata from .utils import ( + affine_to_spacing, compute_importance_map, compute_shape_offset, convert_tables_to_dicts, @@ -63,10 +79,12 @@ create_file_basename, decollate_batch, dense_patch_slices, + get_extra_metadata_keys, get_random_patch, get_valid_patch_size, is_supported_format, iter_patch, + iter_patch_position, iter_patch_slices, json_hashing, list_data_collate, @@ -76,6 +94,8 @@ partition_dataset_classes, pickle_hashing, rectify_header_sform_qform, + remove_extra_metadata, + remove_keys, reorient_spatial_axes, resample_datalist, select_cross_validation_folds, @@ -85,3 +105,29 @@ worker_init_fn, zoom_affine, ) +from .wsi_datasets import MaskedPatchWSIDataset, PatchWSIDataset, SlidingPatchWSIDataset +from .wsi_reader import BaseWSIReader, CuCIMWSIReader, OpenSlideWSIReader, WSIReader + +with contextlib.suppress(BaseException): + from multiprocessing.reduction import ForkingPickler + + def _rebuild_meta(cls, storage, metadata): + storage_offset, size, stride, meta_obj, applied_operations = metadata + t = cls([], meta=meta_obj, applied_operations=applied_operations, dtype=storage.dtype, device=storage.device) + t.set_(storage._untyped() if hasattr(storage, "_untyped") else storage, storage_offset, size, stride) + return t + + def reduce_meta_tensor(meta_tensor): + storage = meta_tensor.storage() + if storage.is_cuda: + raise NotImplementedError("sharing CUDA metatensor across processes not implemented") + metadata = ( + meta_tensor.storage_offset(), + meta_tensor.size(), + meta_tensor.stride(), + meta_tensor.meta, + meta_tensor.applied_operations, + ) + return _rebuild_meta, (type(meta_tensor), storage, metadata) + + ForkingPickler.register(MetaTensor, reduce_meta_tensor) diff --git a/monai/data/box_utils.py b/monai/data/box_utils.py new file mode 100644 index 0000000000..65e5a06c74 --- /dev/null +++ b/monai/data/box_utils.py @@ -0,0 +1,1134 @@ +# 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 utility module mainly supports rectangular bounding boxes with a few +different parameterizations and methods for converting between them. It +provides reliable access to the spatial coordinates of the box vertices in the +"canonical ordering": +[xmin, ymin, xmax, ymax] for 2D and [xmin, ymin, zmin, xmax, ymax, zmax] for 3D. +We currently define this ordering as `monai.data.box_utils.StandardMode` and +the rest of the detection pipelines mainly assumes boxes in `StandardMode`. +""" + +import inspect +import warnings +from abc import ABC, abstractmethod +from copy import deepcopy +from typing import Callable, Dict, Sequence, Tuple, Type, Union + +import numpy as np +import torch + +from monai.config.type_definitions import NdarrayOrTensor +from monai.utils import look_up_option +from monai.utils.enums import BoxModeName +from monai.utils.type_conversion import convert_data_type, convert_to_dst_type + +# We support 2-D or 3-D bounding boxes +SUPPORTED_SPATIAL_DIMS = [2, 3] + + +# TO_REMOVE = 0.0 if the bottom-right corner pixel/voxel is not included in the boxes, +# i.e., when xmin=1., xmax=2., we have w = 1. +# TO_REMOVE = 1.0 if the bottom-right corner pixel/voxel is included in the boxes, +# i.e., when xmin=1., xmax=2., we have w = 2. +# Currently, only `TO_REMOVE = 0.0` is supported +TO_REMOVE = 0.0 # xmax-xmin = w -TO_REMOVE. + +# Some torch functions do not support half precision. +# We therefore compute those functions under COMPUTE_DTYPE +COMPUTE_DTYPE = torch.float32 + + +class BoxMode(ABC): + """ + An abstract class of a ``BoxMode``. + + A ``BoxMode`` is callable that converts box mode of ``boxes``, which are Nx4 (2D) or Nx6 (3D) torch tensor or ndarray. + ``BoxMode`` has several subclasses that represents different box modes, including + + - :class:`~monai.data.box_utils.CornerCornerModeTypeA`: + represents [xmin, ymin, xmax, ymax] for 2D and [xmin, ymin, zmin, xmax, ymax, zmax] for 3D + - :class:`~monai.data.box_utils.CornerCornerModeTypeB`: + represents [xmin, xmax, ymin, ymax] for 2D and [xmin, xmax, ymin, ymax, zmin, zmax] for 3D + - :class:`~monai.data.box_utils.CornerCornerModeTypeC`: + represents [xmin, ymin, xmax, ymax] for 2D and [xmin, ymin, xmax, ymax, zmin, zmax] for 3D + - :class:`~monai.data.box_utils.CornerSizeMode`: + represents [xmin, ymin, xsize, ysize] for 2D and [xmin, ymin, zmin, xsize, ysize, zsize] for 3D + - :class:`~monai.data.box_utils.CenterSizeMode`: + represents [xcenter, ycenter, xsize, ysize] for 2D and [xcenter, ycenter, zcenter, xsize, ysize, zsize] for 3D + + We currently define ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, + and monai detection pipelines mainly assume ``boxes`` are in ``StandardMode``. + + The implementation should be aware of: + + - remember to define class variable ``name``, + a dictionary that maps ``spatial_dims`` to :class:`~monai.utils.enums.BoxModeName`. + - :func:`~monai.data.box_utils.BoxMode.boxes_to_corners` and :func:`~monai.data.box_utils.BoxMode.corners_to_boxes` + should not modify inputs in place. + """ + + # a dictionary that maps spatial_dims to monai.utils.enums.BoxModeName. + name: Dict[int, BoxModeName] = {} + + @classmethod + def get_name(cls, spatial_dims: int) -> str: + """ + Get the mode name for the given spatial dimension using class variable ``name``. + + Args: + spatial_dims: number of spatial dimensions of the bounding boxes. + + Returns: + ``str``: mode string name + """ + return cls.name[spatial_dims].value + + @abstractmethod + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + """ + Convert the bounding boxes of the current mode to corners. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor + + Returns: + ``Tuple``: corners of boxes, 4-element or 6-element tuple, each element is a Nx1 torch tensor. + It represents (xmin, ymin, xmax, ymax) or (xmin, ymin, zmin, xmax, ymax, zmax) + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + boxmode = BoxMode() + boxmode.boxes_to_corners(boxes) # will return a 6-element tuple, each element is a 10x1 tensor + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + @abstractmethod + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + """ + Convert the given box corners to the bounding boxes of the current mode. + + Args: + corners: corners of boxes, 4-element or 6-element tuple, each element is a Nx1 torch tensor. + It represents (xmin, ymin, xmax, ymax) or (xmin, ymin, zmin, xmax, ymax, zmax) + + Returns: + ``Tensor``: bounding boxes, Nx4 or Nx6 torch tensor + + Example: + .. code-block:: python + + corners = (torch.ones(10,1), torch.ones(10,1), torch.ones(10,1), torch.ones(10,1)) + boxmode = BoxMode() + boxmode.corners_to_boxes(corners) # will return a 10x4 tensor + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + +class CornerCornerModeTypeA(BoxMode): + """ + A subclass of ``BoxMode``. + + Also represented as "xyxy" or "xyzxyz", with format of + [xmin, ymin, xmax, ymax] or [xmin, ymin, zmin, xmax, ymax, zmax]. + + Example: + .. code-block:: python + + CornerCornerModeTypeA.get_name(spatial_dims=2) # will return "xyxy" + CornerCornerModeTypeA.get_name(spatial_dims=3) # will return "xyzxyz" + """ + + name = {2: BoxModeName.XYXY, 3: BoxModeName.XYZXYZ} + + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + corners: Tuple + corners = boxes.split(1, dim=-1) + return corners + + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + boxes: torch.Tensor + boxes = torch.cat(tuple(corners), dim=-1) + return boxes + + +class CornerCornerModeTypeB(BoxMode): + """ + A subclass of ``BoxMode``. + + Also represented as "xxyy" or "xxyyzz", with format of + [xmin, xmax, ymin, ymax] or [xmin, xmax, ymin, ymax, zmin, zmax]. + + Example: + .. code-block:: python + + CornerCornerModeTypeB.get_name(spatial_dims=2) # will return "xxyy" + CornerCornerModeTypeB.get_name(spatial_dims=3) # will return "xxyyzz" + """ + + name = {2: BoxModeName.XXYY, 3: BoxModeName.XXYYZZ} + + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + corners: Tuple + spatial_dims = get_spatial_dims(boxes=boxes) + if spatial_dims == 3: + xmin, xmax, ymin, ymax, zmin, zmax = boxes.split(1, dim=-1) + corners = xmin, ymin, zmin, xmax, ymax, zmax + elif spatial_dims == 2: + xmin, xmax, ymin, ymax = boxes.split(1, dim=-1) + corners = xmin, ymin, xmax, ymax + return corners + + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + boxes: torch.Tensor + spatial_dims = get_spatial_dims(corners=corners) + if spatial_dims == 3: + boxes = torch.cat((corners[0], corners[3], corners[1], corners[4], corners[2], corners[5]), dim=-1) + elif spatial_dims == 2: + boxes = torch.cat((corners[0], corners[2], corners[1], corners[3]), dim=-1) + return boxes + + +class CornerCornerModeTypeC(BoxMode): + """ + A subclass of ``BoxMode``. + + Also represented as "xyxy" or "xyxyzz", with format of + [xmin, ymin, xmax, ymax] or [xmin, ymin, xmax, ymax, zmin, zmax]. + + Example: + .. code-block:: python + + CornerCornerModeTypeC.get_name(spatial_dims=2) # will return "xyxy" + CornerCornerModeTypeC.get_name(spatial_dims=3) # will return "xyxyzz" + """ + + name = {2: BoxModeName.XYXY, 3: BoxModeName.XYXYZZ} + + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + corners: Tuple + spatial_dims = get_spatial_dims(boxes=boxes) + if spatial_dims == 3: + xmin, ymin, xmax, ymax, zmin, zmax = boxes.split(1, dim=-1) + corners = xmin, ymin, zmin, xmax, ymax, zmax + elif spatial_dims == 2: + corners = boxes.split(1, dim=-1) + return corners + + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + boxes: torch.Tensor + spatial_dims = get_spatial_dims(corners=corners) + if spatial_dims == 3: + boxes = torch.cat((corners[0], corners[1], corners[3], corners[4], corners[2], corners[5]), dim=-1) + elif spatial_dims == 2: + boxes = torch.cat(tuple(corners), dim=-1) + return boxes + + +class CornerSizeMode(BoxMode): + """ + A subclass of ``BoxMode``. + + Also represented as "xywh" or "xyzwhd", with format of + [xmin, ymin, xsize, ysize] or [xmin, ymin, zmin, xsize, ysize, zsize]. + + Example: + .. code-block:: python + + CornerSizeMode.get_name(spatial_dims=2) # will return "xywh" + CornerSizeMode.get_name(spatial_dims=3) # will return "xyzwhd" + """ + + name = {2: BoxModeName.XYWH, 3: BoxModeName.XYZWHD} + + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + corners: Tuple + # convert to float32 when computing torch.clamp, which does not support float16 + box_dtype = boxes.dtype + + spatial_dims = get_spatial_dims(boxes=boxes) + if spatial_dims == 3: + xmin, ymin, zmin, w, h, d = boxes.split(1, dim=-1) + xmax = xmin + (w - TO_REMOVE).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymax = ymin + (h - TO_REMOVE).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + zmax = zmin + (d - TO_REMOVE).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + corners = xmin, ymin, zmin, xmax, ymax, zmax + elif spatial_dims == 2: + xmin, ymin, w, h = boxes.split(1, dim=-1) + xmax = xmin + (w - TO_REMOVE).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymax = ymin + (h - TO_REMOVE).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + corners = xmin, ymin, xmax, ymax + return corners + + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + boxes: torch.Tensor + spatial_dims = get_spatial_dims(corners=corners) + if spatial_dims == 3: + xmin, ymin, zmin, xmax, ymax, zmax = corners[0], corners[1], corners[2], corners[3], corners[4], corners[5] + boxes = torch.cat( + (xmin, ymin, zmin, xmax - xmin + TO_REMOVE, ymax - ymin + TO_REMOVE, zmax - zmin + TO_REMOVE), dim=-1 + ) + elif spatial_dims == 2: + xmin, ymin, xmax, ymax = corners[0], corners[1], corners[2], corners[3] + boxes = torch.cat((xmin, ymin, xmax - xmin + TO_REMOVE, ymax - ymin + TO_REMOVE), dim=-1) + return boxes + + +class CenterSizeMode(BoxMode): + """ + A subclass of ``BoxMode``. + + Also represented as "ccwh" or "cccwhd", with format of + [xmin, ymin, xsize, ysize] or [xmin, ymin, zmin, xsize, ysize, zsize]. + + Example: + .. code-block:: python + + CenterSizeMode.get_name(spatial_dims=2) # will return "ccwh" + CenterSizeMode.get_name(spatial_dims=3) # will return "cccwhd" + """ + + name = {2: BoxModeName.CCWH, 3: BoxModeName.CCCWHD} + + def boxes_to_corners(self, boxes: torch.Tensor) -> Tuple: + corners: Tuple + # convert to float32 when computing torch.clamp, which does not support float16 + box_dtype = boxes.dtype + + spatial_dims = get_spatial_dims(boxes=boxes) + if spatial_dims == 3: + xc, yc, zc, w, h, d = boxes.split(1, dim=-1) + xmin = xc - ((w - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + xmax = xc + ((w - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymin = yc - ((h - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymax = yc + ((h - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + zmin = zc - ((d - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + zmax = zc + ((d - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + corners = xmin, ymin, zmin, xmax, ymax, zmax + elif spatial_dims == 2: + xc, yc, w, h = boxes.split(1, dim=-1) + xmin = xc - ((w - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + xmax = xc + ((w - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymin = yc - ((h - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + ymax = yc + ((h - TO_REMOVE) / 2.0).to(dtype=COMPUTE_DTYPE).clamp(min=0).to(dtype=box_dtype) + corners = xmin, ymin, xmax, ymax + return corners + + def corners_to_boxes(self, corners: Sequence) -> torch.Tensor: + boxes: torch.Tensor + spatial_dims = get_spatial_dims(corners=corners) + if spatial_dims == 3: + xmin, ymin, zmin, xmax, ymax, zmax = corners[0], corners[1], corners[2], corners[3], corners[4], corners[5] + boxes = torch.cat( + ( + (xmin + xmax + TO_REMOVE) / 2.0, + (ymin + ymax + TO_REMOVE) / 2.0, + (zmin + zmax + TO_REMOVE) / 2.0, + xmax - xmin + TO_REMOVE, + ymax - ymin + TO_REMOVE, + zmax - zmin + TO_REMOVE, + ), + dim=-1, + ) + elif spatial_dims == 2: + xmin, ymin, xmax, ymax = corners[0], corners[1], corners[2], corners[3] + boxes = torch.cat( + ( + (xmin + xmax + TO_REMOVE) / 2.0, + (ymin + ymax + TO_REMOVE) / 2.0, + xmax - xmin + TO_REMOVE, + ymax - ymin + TO_REMOVE, + ), + dim=-1, + ) + return boxes + + +# We support the conversion between several box modes, i.e., representation of a bounding boxes +SUPPORTED_MODES = [CornerCornerModeTypeA, CornerCornerModeTypeB, CornerCornerModeTypeC, CornerSizeMode, CenterSizeMode] +# The standard box mode we use in all the box util functions +StandardMode = CornerCornerModeTypeA + + +def get_spatial_dims( + boxes: Union[torch.Tensor, np.ndarray, None] = None, + points: Union[torch.Tensor, np.ndarray, None] = None, + corners: Union[Sequence, None] = None, + spatial_size: Union[Sequence[int], torch.Tensor, np.ndarray, None] = None, +) -> int: + """ + Get spatial dimension for the giving setting and check the validity of them. + Missing input is allowed. But at least one of the input value should be given. + It raises ValueError if the dimensions of multiple inputs do not match with each other. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray + points: point coordinates, [x, y] or [x, y, z], Nx2 or Nx3 torch tensor or ndarray + corners: corners of boxes, 4-element or 6-element tuple, each element is a Nx1 torch tensor or ndarray + spatial_size: The spatial size of the image where the boxes are attached. + len(spatial_size) should be in [2, 3]. + + Returns: + ``int``: spatial_dims, number of spatial dimensions of the bounding boxes. + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + get_spatial_dims(boxes, spatial_size=[100,200,200]) # will return 3 + get_spatial_dims(boxes, spatial_size=[100,200]) # will raise ValueError + get_spatial_dims(boxes) # will return 3 + """ + spatial_dims_set = set() + + # Check the validity of each input and add its corresponding spatial_dims to spatial_dims_set + if boxes is not None: + if int(boxes.shape[1]) not in [4, 6]: + raise ValueError( + f"Currently we support only boxes with shape [N,4] or [N,6], got boxes with shape {boxes.shape}." + ) + spatial_dims_set.add(int(boxes.shape[1] / 2)) + if points is not None: + if int(points.shape[1]) not in SUPPORTED_SPATIAL_DIMS: + raise ValueError( + f"Currently we support only points with shape [N,2] or [N,3], got boxes with shape {points.shape}." + ) + spatial_dims_set.add(int(points.shape[1])) + if corners is not None: + if len(corners) not in [4, 6]: + raise ValueError( + f"Currently we support only boxes with shape [N,4] or [N,6], got box corner tuple with length {len(corners)}." + ) + spatial_dims_set.add(len(corners) // 2) + if spatial_size is not None: + if len(spatial_size) not in SUPPORTED_SPATIAL_DIMS: + raise ValueError( + f"Currently we support only boxes on 2-D and 3-D images, got image spatial_size {spatial_size}." + ) + spatial_dims_set.add(len(spatial_size)) + + # Get spatial_dims from spatial_dims_set, which contains only unique values + spatial_dims_list = list(spatial_dims_set) + if len(spatial_dims_list) == 0: + raise ValueError("At least one of the inputs needs to be non-empty.") + + if len(spatial_dims_list) == 1: + spatial_dims = int(spatial_dims_list[0]) + spatial_dims = look_up_option(spatial_dims, supported=[2, 3]) + return int(spatial_dims) + + raise ValueError("The dimensions of multiple inputs should match with each other.") + + +def get_boxmode(mode: Union[str, BoxMode, Type[BoxMode], None] = None, *args, **kwargs) -> BoxMode: + """ + This function that return a :class:`~monai.data.box_utils.BoxMode` object giving a representation of box mode + + Args: + mode: a representation of box mode. If it is not given, this func will assume it is ``StandardMode()``. + + Note: + ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, + also represented as "xyxy" for 2D and "xyzxyz" for 3D. + + mode can be: + #. str: choose from :class:`~monai.utils.enums.BoxModeName`, for example, + - "xyxy": boxes has format [xmin, ymin, xmax, ymax] + - "xyzxyz": boxes has format [xmin, ymin, zmin, xmax, ymax, zmax] + - "xxyy": boxes has format [xmin, xmax, ymin, ymax] + - "xxyyzz": boxes has format [xmin, xmax, ymin, ymax, zmin, zmax] + - "xyxyzz": boxes has format [xmin, ymin, xmax, ymax, zmin, zmax] + - "xywh": boxes has format [xmin, ymin, xsize, ysize] + - "xyzwhd": boxes has format [xmin, ymin, zmin, xsize, ysize, zsize] + - "ccwh": boxes has format [xcenter, ycenter, xsize, ysize] + - "cccwhd": boxes has format [xcenter, ycenter, zcenter, xsize, ysize, zsize] + #. BoxMode class: choose from the subclasses of :class:`~monai.data.box_utils.BoxMode`, for example, + - CornerCornerModeTypeA: equivalent to "xyxy" or "xyzxyz" + - CornerCornerModeTypeB: equivalent to "xxyy" or "xxyyzz" + - CornerCornerModeTypeC: equivalent to "xyxy" or "xyxyzz" + - CornerSizeMode: equivalent to "xywh" or "xyzwhd" + - CenterSizeMode: equivalent to "ccwh" or "cccwhd" + #. BoxMode object: choose from the subclasses of :class:`~monai.data.box_utils.BoxMode`, for example, + - CornerCornerModeTypeA(): equivalent to "xyxy" or "xyzxyz" + - CornerCornerModeTypeB(): equivalent to "xxyy" or "xxyyzz" + - CornerCornerModeTypeC(): equivalent to "xyxy" or "xyxyzz" + - CornerSizeMode(): equivalent to "xywh" or "xyzwhd" + - CenterSizeMode(): equivalent to "ccwh" or "cccwhd" + #. None: will assume mode is ``StandardMode()`` + + Returns: + BoxMode object + + Example: + .. code-block:: python + + mode = "xyzxyz" + get_boxmode(mode) # will return CornerCornerModeTypeA() + """ + if isinstance(mode, BoxMode): + return mode + + if inspect.isclass(mode) and issubclass(mode, BoxMode): + return mode(*args, **kwargs) + + if isinstance(mode, str): + for m in SUPPORTED_MODES: + for n in SUPPORTED_SPATIAL_DIMS: + if inspect.isclass(m) and issubclass(m, BoxMode) and m.get_name(n) == mode: + return m(*args, **kwargs) + + if mode is not None: + raise ValueError(f"Unsupported box mode: {mode}.") + return StandardMode(*args, **kwargs) + + +def convert_box_mode( + boxes: NdarrayOrTensor, + src_mode: Union[str, BoxMode, Type[BoxMode], None] = None, + dst_mode: Union[str, BoxMode, Type[BoxMode], None] = None, +) -> NdarrayOrTensor: + """ + This function converts the boxes in src_mode to the dst_mode. + + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. + src_mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``mode`` in :func:`~monai.data.box_utils.get_boxmode`. + dst_mode: target box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``mode`` in :func:`~monai.data.box_utils.get_boxmode`. + + Returns: + bounding boxes with target mode, with same data type as ``boxes``, does not share memory with ``boxes`` + + Example: + .. code-block:: python + + boxes = torch.ones(10,4) + # The following three lines are equivalent + # They convert boxes with format [xmin, ymin, xmax, ymax] to [xcenter, ycenter, xsize, ysize]. + convert_box_mode(boxes=boxes, src_mode="xyxy", dst_mode="ccwh") + convert_box_mode(boxes=boxes, src_mode="xyxy", dst_mode=monai.data.box_utils.CenterSizeMode) + convert_box_mode(boxes=boxes, src_mode="xyxy", dst_mode=monai.data.box_utils.CenterSizeMode()) + """ + src_boxmode = get_boxmode(src_mode) + dst_boxmode = get_boxmode(dst_mode) + + # if mode not changed, deepcopy the original boxes + if isinstance(src_boxmode, type(dst_boxmode)): + return deepcopy(boxes) + + # convert box mode + # convert numpy to tensor if needed + boxes_t, *_ = convert_data_type(boxes, torch.Tensor) + + # convert boxes to corners + corners = src_boxmode.boxes_to_corners(boxes_t) + + # check validity of corners + spatial_dims = get_spatial_dims(boxes=boxes_t) + for axis in range(0, spatial_dims): + if (corners[spatial_dims + axis] < corners[axis]).sum() > 0: + warnings.warn("Given boxes has invalid values. The box size must be non-negative.") + + # convert corners to boxes + boxes_t_dst = dst_boxmode.corners_to_boxes(corners) + + # convert tensor back to numpy if needed + boxes_dst, *_ = convert_to_dst_type(src=boxes_t_dst, dst=boxes) + return boxes_dst + + +def convert_box_to_standard_mode( + boxes: NdarrayOrTensor, mode: Union[str, BoxMode, Type[BoxMode], None] = None +) -> NdarrayOrTensor: + """ + Convert given boxes to standard mode. + Standard mode is "xyxy" or "xyzxyz", + representing box format of [xmin, ymin, xmax, ymax] or [xmin, ymin, zmin, xmax, ymax, zmax]. + + Args: + boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. + mode: source box mode. If it is not given, this func will assume it is ``StandardMode()``. + It follows the same format with ``mode`` in :func:`~monai.data.box_utils.get_boxmode`. + + Returns: + bounding boxes with standard mode, with same data type as ``boxes``, does not share memory with ``boxes`` + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + # The following two lines are equivalent + # They convert boxes with format [xmin, xmax, ymin, ymax, zmin, zmax] to [xmin, ymin, zmin, xmax, ymax, zmax] + convert_box_to_standard_mode(boxes=boxes, mode="xxyyzz") + convert_box_mode(boxes=boxes, src_mode="xxyyzz", dst_mode="xyzxyz") + """ + return convert_box_mode(boxes=boxes, src_mode=mode, dst_mode=StandardMode()) + + +def box_centers(boxes: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Compute center points of boxes + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + center points with size of (N, spatial_dims) + + """ + spatial_dims = get_spatial_dims(boxes=boxes) + return convert_box_mode(boxes=boxes, src_mode=StandardMode, dst_mode=CenterSizeMode)[:, :spatial_dims] + + +def centers_in_boxes(centers: NdarrayOrTensor, boxes: NdarrayOrTensor, eps: float = 0.01) -> NdarrayOrTensor: + """ + Checks which center points are within boxes + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. + centers: center points, Nx2 or Nx3 torch tensor or ndarray. + eps: minimum distance to border of boxes. + + Returns: + boolean array indicating which center points are within the boxes, sized (N,). + + Reference: + https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/ops.py + + """ + spatial_dims = get_spatial_dims(boxes=boxes) + + # compute relative position of centers compared to borders + # should be non-negative if centers are within boxes + center_to_border = [centers[:, axis] - boxes[:, axis] for axis in range(spatial_dims)] + [ + boxes[:, axis + spatial_dims] - centers[:, axis] for axis in range(spatial_dims) + ] + + if isinstance(boxes, np.ndarray): + min_center_to_border: np.ndarray = np.stack(center_to_border, axis=1).min(axis=1) + return min_center_to_border > eps # array[bool] + + return torch.stack(center_to_border, dim=1).to(COMPUTE_DTYPE).min(dim=1)[0] > eps # type: ignore + + +def boxes_center_distance( + boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor, euclidean: bool = True +) -> Tuple[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor]: + """ + Distance of center points between two sets of boxes + + Args: + boxes1: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + boxes2: bounding boxes, Mx4 or Mx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + euclidean: computed the euclidean distance otherwise it uses the l1 distance + + Returns: + - The pairwise distances for every element in boxes1 and boxes2, + with size of (N,M) and same data type as ``boxes1``. + - Center points of boxes1, with size of (N,spatial_dims) and same data type as ``boxes1``. + - Center points of boxes2, with size of (M,spatial_dims) and same data type as ``boxes1``. + + Reference: + https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/ops.py + + """ + + if not isinstance(boxes1, type(boxes2)): + warnings.warn(f"boxes1 is {type(boxes1)}, while boxes2 is {type(boxes2)}. The result will be {type(boxes1)}.") + + # convert numpy to tensor if needed + boxes1_t, *_ = convert_data_type(boxes1, torch.Tensor) + boxes2_t, *_ = convert_data_type(boxes2, torch.Tensor) + + center1 = box_centers(boxes1_t.to(COMPUTE_DTYPE)) # (N, spatial_dims) + center2 = box_centers(boxes2_t.to(COMPUTE_DTYPE)) # (M, spatial_dims) + + if euclidean: + dists = (center1[:, None] - center2[None]).pow(2).sum(-1).sqrt() + else: + # before sum: (N, M, spatial_dims) + dists = (center1[:, None] - center2[None]).sum(-1) + + # convert tensor back to numpy if needed + (dists, center1, center2), *_ = convert_to_dst_type(src=(dists, center1, center2), dst=boxes1) + return dists, center1, center2 + + +def is_valid_box_values(boxes: NdarrayOrTensor) -> bool: + """ + This function checks whether the box size is non-negative. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + whether ``boxes`` is valid + """ + spatial_dims = get_spatial_dims(boxes=boxes) + for axis in range(0, spatial_dims): + if (boxes[:, spatial_dims + axis] < boxes[:, axis]).sum() > 0: + return False + return True + + +def box_area(boxes: NdarrayOrTensor) -> NdarrayOrTensor: + """ + This function computes the area (2D) or volume (3D) of each box. + Half precision is not recommended for this function as it may cause overflow, especially for 3D images. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + area (2D) or volume (3D) of boxes, with size of (N,). + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + # we do computation with torch.float32 to avoid overflow + compute_dtype = torch.float32 + area = box_area(boxes=boxes.to(dtype=compute_dtype)) # torch.float32, size of (10,) + """ + + if not is_valid_box_values(boxes): + raise ValueError("Given boxes has invalid values. The box size must be non-negative.") + + spatial_dims = get_spatial_dims(boxes=boxes) + + area = boxes[:, spatial_dims] - boxes[:, 0] + TO_REMOVE + for axis in range(1, spatial_dims): + area = area * (boxes[:, axis + spatial_dims] - boxes[:, axis] + TO_REMOVE) + + # convert numpy to tensor if needed + area_t, *_ = convert_data_type(area, torch.Tensor) + + # check if NaN or Inf, especially for half precision + if area_t.isnan().any() or area_t.isinf().any(): + if area_t.dtype is torch.float16: + raise ValueError("Box area is NaN or Inf. boxes is float16. Please change to float32 and test it again.") + else: + raise ValueError("Box area is NaN or Inf.") + + return area + + +def _box_inter_union( + boxes1_t: torch.Tensor, boxes2_t: torch.Tensor, compute_dtype: torch.dtype = torch.float32 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This internal function computes the intersection and union area of two set of boxes. + + Args: + boxes1: bounding boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + boxes2: bounding boxes, Mx4 or Mx6 torch tensor. The box mode is assumed to be ``StandardMode`` + compute_dtype: default torch.float32, dtype with which the results will be computed + + Returns: + inter, with size of (N,M) and dtype of ``compute_dtype``. + union, with size of (N,M) and dtype of ``compute_dtype``. + + """ + spatial_dims = get_spatial_dims(boxes=boxes1_t) + + # compute area with float32 + area1 = box_area(boxes=boxes1_t.to(dtype=compute_dtype)) # (N,) + area2 = box_area(boxes=boxes2_t.to(dtype=compute_dtype)) # (M,) + + # get the left top and right bottom points for the NxM combinations + lt = torch.max(boxes1_t[:, None, :spatial_dims], boxes2_t[:, :spatial_dims]).to( + dtype=compute_dtype + ) # (N,M,spatial_dims) left top + rb = torch.min(boxes1_t[:, None, spatial_dims:], boxes2_t[:, spatial_dims:]).to( + dtype=compute_dtype + ) # (N,M,spatial_dims) right bottom + + # compute size for the intersection region for the NxM combinations + wh = (rb - lt + TO_REMOVE).clamp(min=0) # (N,M,spatial_dims) + inter = torch.prod(wh, dim=-1, keepdim=False) # (N,M) + + union = area1[:, None] + area2 - inter + return inter, union + + +def box_iou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Compute the intersection over union (IoU) of two set of boxes. + + Args: + boxes1: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + boxes2: bounding boxes, Mx4 or Mx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + IoU, with size of (N,M) and same data type as ``boxes1`` + + """ + + if not isinstance(boxes1, type(boxes2)): + warnings.warn(f"boxes1 is {type(boxes1)}, while boxes2 is {type(boxes2)}. The result will be {type(boxes1)}.") + + # convert numpy to tensor if needed + boxes1_t, *_ = convert_data_type(boxes1, torch.Tensor) + boxes2_t, *_ = convert_data_type(boxes2, torch.Tensor) + + # we do computation with compute_dtype to avoid overflow + box_dtype = boxes1_t.dtype + + inter, union = _box_inter_union(boxes1_t, boxes2_t, compute_dtype=COMPUTE_DTYPE) + + # compute IoU and convert back to original box_dtype + iou_t = inter / (union + torch.finfo(COMPUTE_DTYPE).eps) # (N,M) + iou_t = iou_t.to(dtype=box_dtype) + + # check if NaN or Inf + if torch.isnan(iou_t).any() or torch.isinf(iou_t).any(): + raise ValueError("Box IoU is NaN or Inf.") + + # convert tensor back to numpy if needed + iou, *_ = convert_to_dst_type(src=iou_t, dst=boxes1) + return iou + + +def box_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Compute the generalized intersection over union (GIoU) of two sets of boxes. + The two inputs can have different shapes and the func return an NxM matrix, + (in contrary to :func:`~monai.data.box_utils.box_pair_giou` , which requires the inputs to have the same + shape and returns ``N`` values). + + Args: + boxes1: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + boxes2: bounding boxes, Mx4 or Mx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + + Returns: + GIoU, with size of (N,M) and same data type as ``boxes1`` + + Reference: + https://giou.stanford.edu/GIoU.pdf + + """ + + if not isinstance(boxes1, type(boxes2)): + warnings.warn(f"boxes1 is {type(boxes1)}, while boxes2 is {type(boxes2)}. The result will be {type(boxes1)}.") + + # convert numpy to tensor if needed + boxes1_t, *_ = convert_data_type(boxes1, torch.Tensor) + boxes2_t, *_ = convert_data_type(boxes2, torch.Tensor) + + spatial_dims = get_spatial_dims(boxes=boxes1_t) + + # we do computation with compute_dtype to avoid overflow + box_dtype = boxes1_t.dtype + + inter, union = _box_inter_union(boxes1_t, boxes2_t, compute_dtype=COMPUTE_DTYPE) + iou = inter / (union + torch.finfo(COMPUTE_DTYPE).eps) # (N,M) + + # Enclosure + # get the left top and right bottom points for the NxM combinations + lt = torch.min(boxes1_t[:, None, :spatial_dims], boxes2_t[:, :spatial_dims]).to( + dtype=COMPUTE_DTYPE + ) # (N,M,spatial_dims) left top + rb = torch.max(boxes1_t[:, None, spatial_dims:], boxes2_t[:, spatial_dims:]).to( + dtype=COMPUTE_DTYPE + ) # (N,M,spatial_dims) right bottom + + # compute size for the enclosure region for the NxM combinations + wh = (rb - lt + TO_REMOVE).clamp(min=0) # (N,M,spatial_dims) + enclosure = torch.prod(wh, dim=-1, keepdim=False) # (N,M) + + # GIoU + giou_t = iou - (enclosure - union) / (enclosure + torch.finfo(COMPUTE_DTYPE).eps) + giou_t = giou_t.to(dtype=box_dtype) + if torch.isnan(giou_t).any() or torch.isinf(giou_t).any(): + raise ValueError("Box GIoU is NaN or Inf.") + + # convert tensor back to numpy if needed + giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1) + return giou + + +def box_pair_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTensor: + """ + Compute the generalized intersection over union (GIoU) of a pair of boxes. + The two inputs should have the same shape and the func return an (N,) array, + (in contrary to :func:`~monai.data.box_utils.box_giou` , which does not require the inputs to have the same + shape and returns ``NxM`` matrix). + + Args: + boxes1: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + boxes2: bounding boxes, same shape with boxes1. The box mode is assumed to be ``StandardMode`` + + Returns: + paired GIoU, with size of (N,) and same data type as ``boxes1`` + + Reference: + https://giou.stanford.edu/GIoU.pdf + + """ + + if not isinstance(boxes1, type(boxes2)): + warnings.warn(f"boxes1 is {type(boxes1)}, while boxes2 is {type(boxes2)}. The result will be {type(boxes1)}.") + + # convert numpy to tensor if needed + boxes1_t, *_ = convert_data_type(boxes1, torch.Tensor) + boxes2_t, *_ = convert_data_type(boxes2, torch.Tensor) + + if boxes1_t.shape != boxes2_t.shape: + raise ValueError("boxes1 and boxes2 should be paired and have same shape.") + + spatial_dims = get_spatial_dims(boxes=boxes1_t) + + # we do computation with compute_dtype to avoid overflow + box_dtype = boxes1_t.dtype + + # compute area + area1 = box_area(boxes=boxes1_t.to(dtype=COMPUTE_DTYPE)) # (N,) + area2 = box_area(boxes=boxes2_t.to(dtype=COMPUTE_DTYPE)) # (N,) + + # Intersection + # get the left top and right bottom points for the boxes pair + lt = torch.max(boxes1_t[:, :spatial_dims], boxes2_t[:, :spatial_dims]).to( + dtype=COMPUTE_DTYPE + ) # (N,spatial_dims) left top + rb = torch.min(boxes1_t[:, spatial_dims:], boxes2_t[:, spatial_dims:]).to( + dtype=COMPUTE_DTYPE + ) # (N,spatial_dims) right bottom + + # compute size for the intersection region for the boxes pair + wh = (rb - lt + TO_REMOVE).clamp(min=0) # (N,spatial_dims) + inter = torch.prod(wh, dim=-1, keepdim=False) # (N,) + + # compute IoU and convert back to original box_dtype + union = area1 + area2 - inter + iou = inter / (union + torch.finfo(COMPUTE_DTYPE).eps) # (N,) + + # Enclosure + # get the left top and right bottom points for the boxes pair + lt = torch.min(boxes1_t[:, :spatial_dims], boxes2_t[:, :spatial_dims]).to( + dtype=COMPUTE_DTYPE + ) # (N,spatial_dims) left top + rb = torch.max(boxes1_t[:, spatial_dims:], boxes2_t[:, spatial_dims:]).to( + dtype=COMPUTE_DTYPE + ) # (N,spatial_dims) right bottom + + # compute size for the enclose region for the boxes pair + wh = (rb - lt + TO_REMOVE).clamp(min=0) # (N,spatial_dims) + enclosure = torch.prod(wh, dim=-1, keepdim=False) # (N,) + + giou_t = iou - (enclosure - union) / (enclosure + torch.finfo(COMPUTE_DTYPE).eps) + giou_t = giou_t.to(dtype=box_dtype) # (N,spatial_dims) + if torch.isnan(giou_t).any() or torch.isinf(giou_t).any(): + raise ValueError("Box GIoU is NaN or Inf.") + + # convert tensor back to numpy if needed + giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1) + return giou + + +def spatial_crop_boxes( + boxes: NdarrayOrTensor, + roi_start: Union[Sequence[int], NdarrayOrTensor], + roi_end: Union[Sequence[int], NdarrayOrTensor], + remove_empty: bool = True, +) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + """ + This function generate the new boxes when the corresponding image is cropped to the given ROI. + When ``remove_empty=True``, it makes sure the bounding boxes are within the new cropped image. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + roi_start: voxel coordinates for start of the crop ROI, negative values allowed. + roi_end: voxel coordinates for end of the crop ROI, negative values allowed. + remove_empty: whether to remove the boxes that are actually empty + + Returns: + - cropped boxes, boxes[keep], does not share memory with original boxes + - ``keep``, it indicates whether each box in ``boxes`` are kept when ``remove_empty=True``. + """ + + # convert numpy to tensor if needed + boxes_t, *_ = convert_data_type(deepcopy(boxes), torch.Tensor) + + # 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_end_t = torch.maximum(roi_end_t, roi_start_t) + + # makes sure the bounding boxes are within the patch + spatial_dims = get_spatial_dims(boxes=boxes, spatial_size=roi_end) + for axis in range(0, spatial_dims): + boxes_t[:, axis].clamp_(min=roi_start_t[axis], max=roi_end_t[axis] - TO_REMOVE) + boxes_t[:, axis + spatial_dims].clamp_(min=roi_start_t[axis], max=roi_end_t[axis] - TO_REMOVE) + boxes_t[:, axis] -= roi_start_t[axis] + boxes_t[:, axis + spatial_dims] -= roi_start_t[axis] + + # remove the boxes that are actually empty + if remove_empty: + keep_t = boxes_t[:, spatial_dims] >= boxes_t[:, 0] + 1 - TO_REMOVE + for axis in range(1, spatial_dims): + keep_t = keep_t & (boxes_t[:, axis + spatial_dims] >= boxes_t[:, axis] + 1 - TO_REMOVE) + boxes_t = boxes_t[keep_t] + else: + keep_t = torch.full_like(boxes_t[:, 0], fill_value=True, dtype=torch.bool) + + # convert tensor back to numpy if needed + boxes_keep, *_ = convert_to_dst_type(src=boxes_t, dst=boxes) + keep, *_ = convert_to_dst_type(src=keep_t, dst=boxes, dtype=keep_t.dtype) + + return boxes_keep, keep + + +def clip_boxes_to_image( + boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], NdarrayOrTensor], remove_empty: bool = True +) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + """ + This function clips the ``boxes`` to makes sure the bounding boxes are within the image. + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: The spatial size of the image where the boxes are attached. len(spatial_size) should be in [2, 3]. + remove_empty: whether to remove the boxes that are actually empty + + Returns: + - clipped boxes, boxes[keep], does not share memory with original boxes + - ``keep``, it indicates whether each box in ``boxes`` are kept when ``remove_empty=True``. + """ + spatial_dims = get_spatial_dims(boxes=boxes, spatial_size=spatial_size) + return spatial_crop_boxes(boxes, roi_start=[0] * spatial_dims, roi_end=spatial_size, remove_empty=remove_empty) + + +def non_max_suppression( + boxes: NdarrayOrTensor, + scores: NdarrayOrTensor, + nms_thresh: float, + max_proposals: int = -1, + box_overlap_metric: Callable = box_iou, +) -> NdarrayOrTensor: + """ + Non-maximum suppression (NMS). + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + scores: prediction scores of the boxes, sized (N,). This function keeps boxes with higher scores. + nms_thresh: threshold of NMS. Discards all overlapping boxes with box_overlap > nms_thresh. + max_proposals: maximum number of boxes it keeps. + If ``max_proposals`` = -1, there is no limit on the number of boxes that are kept. + box_overlap_metric: the metric to compute overlap between boxes. + + Returns: + Indexes of ``boxes`` that are kept after NMS. + + Example: + .. code-block:: python + + boxes = torch.ones(10,6) + scores = torch.ones(10) + keep = non_max_suppression(boxes, scores, num_thresh=0.1) + boxes_after_nms = boxes[keep] + """ + + # returns empty array if boxes is empty + if boxes.shape[0] == 0: + return convert_to_dst_type(src=np.array([]), dst=boxes, dtype=torch.long)[0] + + if boxes.shape[0] != scores.shape[0]: + raise ValueError( + f"boxes and scores should have same length, got boxes shape {boxes.shape}, scores shape {scores.shape}" + ) + + # convert numpy to tensor if needed + boxes_t, *_ = convert_data_type(boxes, torch.Tensor) + scores_t, *_ = convert_to_dst_type(scores, boxes_t) + + # sort boxes in desending order according to the scores + sort_idxs = torch.argsort(scores_t, dim=0, descending=True) + boxes_sort = deepcopy(boxes_t)[sort_idxs, :] + + # initialize the list of picked indexes + pick = [] + idxs = torch.Tensor(list(range(0, boxes_sort.shape[0]))).to(torch.long) + + # keep looping while some indexes still remain in the indexes list + while len(idxs) > 0: + # pick the first index in the indexes list and add the index value to the list of picked indexes + i = int(idxs[0].item()) + pick.append(i) + if len(pick) >= max_proposals >= 1: + break + + # compute the IoU between the rest of the boxes and the box just picked + box_overlap = box_overlap_metric(boxes_sort[idxs, :], boxes_sort[i : i + 1, :]) + + # keep only indexes from the index list that have overlap < nms_thresh + to_keep_idx = (box_overlap <= nms_thresh).flatten() + to_keep_idx[0] = False # always remove idxs[0] + idxs = idxs[to_keep_idx] + + # return only the bounding boxes that were picked using the integer data type + pick_idx = sort_idxs[pick] + + # convert tensor back to numpy if needed + return convert_to_dst_type(src=pick_idx, dst=boxes, dtype=pick_idx.dtype)[0] + + +def batched_nms( + boxes: NdarrayOrTensor, + scores: NdarrayOrTensor, + labels: NdarrayOrTensor, + nms_thresh: float, + max_proposals: int = -1, + box_overlap_metric: Callable = box_iou, +) -> NdarrayOrTensor: + """ + Performs non-maximum suppression in a batched fashion. + Each labels value correspond to a category, and NMS will not be applied between elements of different categories. + + Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/nms.py + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + scores: prediction scores of the boxes, sized (N,). This function keeps boxes with higher scores. + labels: indices of the categories for each one of the boxes. sized(N,), value range is (0, num_classes) + nms_thresh: threshold of NMS. Discards all overlapping boxes with box_overlap > nms_thresh. + max_proposals: maximum number of boxes it keeps. + If ``max_proposals`` = -1, there is no limit on the number of boxes that are kept. + box_overlap_metric: the metric to compute overlap between boxes. + + Returns: + Indexes of ``boxes`` that are kept after NMS. + """ + # returns empty array if boxes is empty + if boxes.shape[0] == 0: + return convert_to_dst_type(src=np.array([]), dst=boxes, dtype=torch.long)[0] + + # convert numpy to tensor if needed + boxes_t, *_ = convert_data_type(boxes, torch.Tensor, dtype=torch.float32) + scores_t, *_ = convert_to_dst_type(scores, boxes_t) + labels_t, *_ = convert_to_dst_type(labels, boxes_t, dtype=torch.long) + + # strategy: in order to perform NMS independently per class. + # we add an offset to all the boxes. The offset is dependent + # only on the class idx, and is large enough so that boxes + # 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] + keep = non_max_suppression(boxes_for_nms, scores_t, nms_thresh, max_proposals, box_overlap_metric) + + # convert tensor back to numpy if needed + return convert_to_dst_type(src=keep, dst=boxes, dtype=keep.dtype)[0] diff --git a/monai/data/csv_saver.py b/monai/data/csv_saver.py index e938bdabf8..36d159d4be 100644 --- a/monai/data/csv_saver.py +++ b/monai/data/csv_saver.py @@ -27,7 +27,7 @@ class CSVSaver: Save the data in a dictionary format cache, and write to a CSV file finally. Typically, the data can be classification predictions, call `save` for single data or call `save_batch` to save a batch of data together, and call `finalize` to write - the cached data into CSV file. If no meta data provided, use index from 0 to save data. + the cached data into CSV file. If no metadata provided, use index from 0 to save data. Note that this saver can't support multi-processing because it reads / writes single CSV file and can't guarantee the data order in multi-processing situation. @@ -88,7 +88,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] Args: data: target data content that save into cache. - meta_data: the meta data information corresponding to the data. + meta_data: the metadata information corresponding to the data. """ save_key = meta_data[Key.FILENAME_OR_OBJ] if meta_data else str(self._data_index) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 3c1fc0abed..59a33bbf20 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -669,7 +669,7 @@ class CacheDataset(Dataset): def __init__( self, data: Sequence, - transform: Union[Sequence[Callable], Callable], + transform: Optional[Union[Sequence[Callable], Callable]] = None, cache_num: int = sys.maxsize, cache_rate: float = 1.0, num_workers: Optional[int] = 1, @@ -847,8 +847,9 @@ class SmartCacheDataset(Randomizable, CacheDataset): Note: This replacement will not work for below cases: 1. Set the `multiprocessing_context` of DataLoader to `spawn`. - 2. Run on windows(the default multiprocessing method is `spawn`) with `num_workers` greater than 0. - 3. Set the `persistent_workers` of DataLoader to `True` with `num_workers` greater than 0. + 2. Launch distributed data parallel with `torch.multiprocessing.spawn`. + 3. Run on windows(the default multiprocessing method is `spawn`) with `num_workers` greater than 0. + 4. Set the `persistent_workers` of DataLoader to `True` with `num_workers` greater than 0. If using MONAI workflows, please add `SmartCacheHandler` to the handler list of trainer, otherwise, please make sure to call `start()`, `update_cache()`, `shutdown()` during training. @@ -856,7 +857,7 @@ class SmartCacheDataset(Randomizable, CacheDataset): Args: data: input data to load and transform to generate dataset for model. transform: transforms to execute operations on input data. - replace_rate: percentage of the cached items to be replaced in every epoch. + replace_rate: percentage of the cached items to be replaced in every epoch (default to 0.1). cache_num: number of items to be cached. Default is `sys.maxsize`. will take the minimum of (cache_num, data_length x cache_rate, data_length). cache_rate: percentage of cached data in total, default is 1.0 (cache all). @@ -883,8 +884,8 @@ class SmartCacheDataset(Randomizable, CacheDataset): def __init__( self, data: Sequence, - transform: Union[Sequence[Callable], Callable], - replace_rate: float, + transform: Optional[Union[Sequence[Callable], Callable]] = None, + replace_rate: float = 0.1, cache_num: int = sys.maxsize, cache_rate: float = 1.0, num_init_workers: Optional[int] = 1, diff --git a/monai/data/dataset_summary.py b/monai/data/dataset_summary.py index b447585d3e..6af46d11ea 100644 --- a/monai/data/dataset_summary.py +++ b/monai/data/dataset_summary.py @@ -18,9 +18,9 @@ from monai.config import KeysCollection from monai.data.dataloader import DataLoader from monai.data.dataset import Dataset +from monai.data.utils import affine_to_spacing from monai.transforms import concatenate -from monai.utils import convert_data_type -from monai.utils.enums import PostFix +from monai.utils import PostFix, convert_data_type DEFAULT_POST_FIX = PostFix.meta() @@ -54,12 +54,12 @@ def __init__( dataset: dataset from which to load the data. image_key: key name of images (default: ``image``). label_key: key name of labels (default: ``label``). - meta_key: explicitly indicate the key of the corresponding meta data dictionary. + meta_key: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, affine, original_shape, etc. + the metadata is a dictionary object which contains: filename, affine, original_shape, etc. if None, will try to construct meta_keys by `{image_key}_{meta_key_postfix}`. - meta_key_postfix: use `{image_key}_{meta_key_postfix}` to fetch the meta data from dict, - the meta data is a dictionary object (default: ``meta_dict``). + meta_key_postfix: use `{image_key}_{meta_key_postfix}` to fetch the metadata from dict, + the metadata is a dictionary object (default: ``meta_dict``). num_workers: how many subprocesses to use for data loading. ``0`` means that the data will be loaded in the main process (default: ``0``). kwargs: other parameters (except `batch_size` and `num_workers`) for DataLoader, @@ -76,15 +76,15 @@ def __init__( def collect_meta_data(self): """ - This function is used to collect the meta data for all images of the dataset. + This function is used to collect the metadata for all images of the dataset. """ for data in self.data_loader: if self.meta_key not in data: - raise ValueError(f"To collect meta data for the dataset, key `{self.meta_key}` must exist in `data`.") + raise ValueError(f"To collect metadata for the dataset, key `{self.meta_key}` must exist in `data`.") self.all_meta_data.append(data[self.meta_key]) - def get_target_spacing(self, spacing_key: str = "pixdim", anisotropic_threshold: int = 3, percentile: float = 10.0): + def get_target_spacing(self, spacing_key: str = "affine", anisotropic_threshold: int = 3, percentile: float = 10.0): """ Calculate the target spacing according to all spacings. If the target spacing is very anisotropic, @@ -93,7 +93,7 @@ def get_target_spacing(self, spacing_key: str = "pixdim", anisotropic_threshold: After loading with `monai.DataLoader`, "pixdim" is in the form of `torch.Tensor` with size `(batch_size, 8)`. Args: - spacing_key: key of spacing in meta data (default: ``pixdim``). + spacing_key: key of the affine used to compute spacing in metadata (default: ``affine``). anisotropic_threshold: threshold to decide if the target spacing is anisotropic (default: ``3``). percentile: for anisotropic target spacing, use the percentile of all spacings of the anisotropic axis to replace that axis. @@ -103,7 +103,8 @@ def get_target_spacing(self, spacing_key: str = "pixdim", anisotropic_threshold: self.collect_meta_data() if spacing_key not in self.all_meta_data[0]: raise ValueError("The provided spacing_key is not in self.all_meta_data.") - all_spacings = concatenate(to_cat=[data[spacing_key][:, 1:4] for data in self.all_meta_data], axis=0) + spacings = [affine_to_spacing(data[spacing_key][0], 3)[None] for data in self.all_meta_data] + all_spacings = concatenate(to_cat=spacings, axis=0) all_spacings, *_ = convert_data_type(data=all_spacings, output_type=np.ndarray, wrap_sequence=True) target_spacing = np.median(all_spacings, axis=0) diff --git a/monai/data/decathlon_datalist.py b/monai/data/decathlon_datalist.py index d2a9c3d220..c1bbabceb9 100644 --- a/monai/data/decathlon_datalist.py +++ b/monai/data/decathlon_datalist.py @@ -122,7 +122,8 @@ def load_decathlon_datalist( if data_list_key not in json_data: raise ValueError(f'Data list {data_list_key} not specified in "{data_list_file_path}".') expected_data = json_data[data_list_key] - if data_list_key == "test": + if data_list_key == "test" and not isinstance(expected_data[0], dict): + # decathlon datalist may save the test images in a list directly instead of dict expected_data = [{"image": i} for i in expected_data] if base_dir is None: diff --git a/monai/data/fft_utils.py b/monai/data/fft_utils.py new file mode 100644 index 0000000000..2d38beed5e --- /dev/null +++ b/monai/data/fft_utils.py @@ -0,0 +1,94 @@ +# 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. + +import torch + +from monai.config.type_definitions import NdarrayOrTensor +from monai.networks.blocks.fft_utils_t import fftn_centered_t, ifftn_centered_t +from monai.utils.type_conversion import convert_data_type, convert_to_dst_type + + +def ifftn_centered(ksp: NdarrayOrTensor, spatial_dims: int, is_complex: bool = True) -> NdarrayOrTensor: + """ + Pytorch-based ifft for spatial_dims-dim signals. "centered" means this function automatically takes care + of the required ifft and fft shifts. This function calls monai.metworks.blocks.fft_utils_t.ifftn_centered_t. + This is equivalent to do fft in numpy based on numpy.fft.ifftn, numpy.fft.fftshift, and numpy.fft.ifftshift + + Args: + ksp: k-space data that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) + is_complex: if True, then the last dimension of the input ksp is expected to be 2 (representing real and imaginary channels) + + Returns: + "out" which is the output image (inverse fourier of ksp) + + Example: + + .. code-block:: python + + import torch + ksp = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts + # output1 and output2 will be identical + output1 = torch.fft.ifftn(torch.view_as_complex(torch.fft.ifftshift(ksp,dim=(-3,-2))), dim=(-2,-1), norm="ortho") + output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) + + output2 = ifftn_centered(ksp, spatial_dims=2, is_complex=True) + """ + # handle numpy format + ksp_t, *_ = convert_data_type(ksp, torch.Tensor) + + # compute ifftn + out_t = ifftn_centered_t(ksp_t, spatial_dims=spatial_dims, is_complex=is_complex) + + # handle numpy format + out, *_ = convert_to_dst_type(src=out_t, dst=ksp) + return out + + +def fftn_centered(im: NdarrayOrTensor, spatial_dims: int, is_complex: bool = True) -> NdarrayOrTensor: + """ + Pytorch-based fft for spatial_dims-dim signals. "centered" means this function automatically takes care + of the required ifft and fft shifts. This function calls monai.metworks.blocks.fft_utils_t.fftn_centered_t. + This is equivalent to do ifft in numpy based on numpy.fft.fftn, numpy.fft.fftshift, and numpy.fft.ifftshift + + Args: + im: image that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) + is_complex: if True, then the last dimension of the input im is expected to be 2 (representing real and imaginary channels) + + Returns: + "out" which is the output kspace (fourier of im) + + Example: + + .. code-block:: python + + import torch + im = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts + # output1 and output2 will be identical + output1 = torch.fft.fftn(torch.view_as_complex(torch.fft.ifftshift(im,dim=(-3,-2))), dim=(-2,-1), norm="ortho") + output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) + + output2 = fftn_centered(im, spatial_dims=2, is_complex=True) + """ + # handle numpy format + im_t, *_ = convert_data_type(im, torch.Tensor) + + # compute ifftn + out_t = fftn_centered_t(im_t, spatial_dims=spatial_dims, is_complex=is_complex) + + # handle numpy format + out, *_ = convert_to_dst_type(src=out_t, dst=im) + return out diff --git a/monai/data/grid_dataset.py b/monai/data/grid_dataset.py index 9eb84a58c9..2b28949419 100644 --- a/monai/data/grid_dataset.py +++ b/monai/data/grid_dataset.py @@ -9,29 +9,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, Iterable, Optional, Sequence, Union +from copy import deepcopy +from typing import Callable, Dict, Hashable, Iterable, Mapping, Optional, Sequence, Union +import numpy as np + +from monai.config import KeysCollection from monai.data.dataset import Dataset from monai.data.iterable_dataset import IterableDataset from monai.data.utils import iter_patch from monai.transforms import apply_transform -from monai.utils import NumpyPadMode, deprecated_arg, ensure_tuple, look_up_option +from monai.utils import NumpyPadMode, deprecated_arg, ensure_tuple, first, look_up_option -__all__ = ["PatchDataset", "GridPatchDataset", "PatchIter"] +__all__ = ["PatchDataset", "GridPatchDataset", "PatchIter", "PatchIterd"] class PatchIter: """ - A class to return a patch generator with predefined properties such as `patch_size`. + Return a patch generator with predefined properties such as `patch_size`. Typically used with :py:class:`monai.data.GridPatchDataset`. + """ def __init__( - self, - patch_size: Sequence[int], - start_pos: Sequence[int] = (), - mode: Union[NumpyPadMode, str] = NumpyPadMode.WRAP, - **pad_opts: Dict, + self, patch_size: Sequence[int], start_pos: Sequence[int] = (), mode: str = NumpyPadMode.WRAP, **pad_opts: Dict ): """ @@ -42,7 +43,8 @@ def __init__( ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} One of the listed string values or a user supplied function. Defaults to ``"wrap"``. See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html - pad_opts: padding options, see numpy.pad + pad_opts: other arguments for the `np.pad` function. + note that `np.pad` treats channel dimension as the first dimension. Note: The `patch_size` is the size of the @@ -52,37 +54,90 @@ def __init__( specified by a `patch_size` of (10, 10, 10). """ - self.patch_size = (None,) + tuple(patch_size) + self.patch_size = (None,) + tuple(patch_size) # expand to have the channel dim self.start_pos = ensure_tuple(start_pos) self.mode: NumpyPadMode = look_up_option(mode, NumpyPadMode) self.pad_opts = pad_opts - def __call__(self, array): + def __call__(self, array: np.ndarray): """ Args: array: the image to generate patches from. + """ yield from iter_patch( array, - patch_size=self.patch_size, # expand to have the channel dim + patch_size=self.patch_size, # type: ignore start_pos=self.start_pos, + overlap=0.0, copy_back=False, mode=self.mode, **self.pad_opts, ) +class PatchIterd: + """ + Dictionary-based wrapper of :py:class:`monai.data.PatchIter`. + Return a patch generator for dictionary data and the coordinate, Typically used + with :py:class:`monai.data.GridPatchDataset`. + Suppose all the expected fields specified by `keys` have same shape. + + Args: + keys: keys of the corresponding items to iterate patches. + patch_size: size of patches to generate slices for, 0/None selects whole dimension + start_pos: starting position in the array, default is 0 for each dimension + mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, + ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + One of the listed string values or a user supplied function. Defaults to ``"wrap"``. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + pad_opts: other arguments for the `np.pad` function. + note that `np.pad` treats channel dimension as the first dimension. + + """ + + coords_key = "patch_coords" + original_spatial_shape_key = "original_spatial_shape" + start_pos_key = "start_pos" + + def __init__( + self, + keys: KeysCollection, + patch_size: Sequence[int], + start_pos: Sequence[int] = (), + mode: str = NumpyPadMode.WRAP, + **pad_opts, + ): + self.keys = ensure_tuple(keys) + self.patch_iter = PatchIter(patch_size=patch_size, start_pos=start_pos, mode=mode, **pad_opts) + + def __call__(self, data: Mapping[Hashable, np.ndarray]): + d = dict(data) + original_spatial_shape = d[first(self.keys)].shape[1:] + + for patch in zip(*[self.patch_iter(d[key]) for key in self.keys]): + coords = patch[0][1] # use the coordinate of the first item + ret = {k: v[0] for k, v in zip(self.keys, patch)} + # fill in the extra keys with unmodified data + for k in set(d.keys()).difference(set(self.keys)): + ret[k] = deepcopy(d[k]) + # also store the `coordinate`, `spatial shape of original image`, `start position` in the dictionary + ret[self.coords_key] = coords + ret[self.original_spatial_shape_key] = original_spatial_shape + ret[self.start_pos_key] = self.patch_iter.start_pos + yield ret, coords + + class GridPatchDataset(IterableDataset): """ - Yields patches from images read from an image dataset. - Typically used with `PatchIter` so that the patches are chosen in a contiguous grid sampling scheme. + Yields patches from data read from an image dataset. + Typically used with `PatchIter` or `PatchIterd` so that the patches are chosen in a contiguous grid sampling scheme. .. code-block:: python import numpy as np - from monai.data import GridPatchDataset, DataLoader, PatchIter - from monai.transforms import RandShiftIntensity + from monai.data import GridPatchDataset, DataLoader, PatchIter, RandShiftIntensity # image-level dataset images = [np.arange(16, dtype=float).reshape(1, 4, 4), @@ -109,7 +164,7 @@ class GridPatchDataset(IterableDataset): data: the data source to read image data from. patch_iter: converts an input image (item from dataset) into a iterable of image patches. `patch_iter(dataset[idx])` must yield a tuple: (patches, coordinates). - see also: :py:class:`monai.data.PatchIter`. + see also: :py:class:`monai.data.PatchIter` or :py:class:`monai.data.PatchIterd`. transform: a callable data transform operates on the patches. with_coordinates: whether to yield the coordinates of each patch, default to `True`. @@ -128,23 +183,19 @@ def __init__( ) -> None: super().__init__(data=data, transform=None) self.patch_iter = patch_iter - self.transform = transform + self.patch_transform = transform self.with_coordinates = with_coordinates def __iter__(self): for image in super().__iter__(): - if not self.with_coordinates: - for patch, *_ in self.patch_iter(image): # patch_iter to yield at least 1 item: patch - out_patch = ( - patch if self.transform is None else apply_transform(self.transform, patch, map_items=False) - ) + for patch, *others in self.patch_iter(image): + out_patch = patch + if self.patch_transform is not None: + out_patch = apply_transform(self.patch_transform, patch, map_items=False) + if self.with_coordinates and len(others) > 0: # patch_iter to yield at least 2 items: patch, coords + yield out_patch, others[0] + else: yield out_patch - else: - for patch, slices, *_ in self.patch_iter(image): # patch_iter to yield at least 2 items: patch, coords - out_patch = ( - patch if self.transform is None else apply_transform(self.transform, patch, map_items=False) - ) - yield out_patch, slices class PatchDataset(Dataset): diff --git a/monai/data/image_dataset.py b/monai/data/image_dataset.py index 51f4e04959..89694a4bb8 100644 --- a/monai/data/image_dataset.py +++ b/monai/data/image_dataset.py @@ -59,7 +59,7 @@ def __init__( image_only: if True return only the image volume, otherwise, return image volume and the metadata. transform_with_metadata: if True, the metadata will be passed to the transforms whenever possible. dtype: if not None convert the loaded image to this data type. - reader: register reader to load image file and meta data, if None, will use the default readers. + reader: register reader to load image file and metadata, if None, will use the default readers. If a string of reader name provided, will construct a reader object with the `*args` and `**kwargs` parameters, supported reader name: "NibabelReader", "PILReader", "ITKReader", "NumpyReader" args: additional parameters for reader if providing a reader name. diff --git a/monai/data/image_reader.py b/monai/data/image_reader.py index 502c6fb93b..743896c99a 100644 --- a/monai/data/image_reader.py +++ b/monai/data/image_reader.py @@ -9,8 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import glob +import os import warnings from abc import ABC, abstractmethod +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple, Union @@ -18,28 +21,46 @@ from torch.utils.data._utils.collate import np_str_obj_array_pattern from monai.config import DtypeLike, KeysCollection, PathLike -from monai.data.utils import correct_nifti_header_if_necessary, is_supported_format, orientation_ras_lps +from monai.data.utils import ( + affine_to_spacing, + correct_nifti_header_if_necessary, + is_supported_format, + orientation_ras_lps, +) from monai.transforms.utility.array import EnsureChannelFirst from monai.utils import ensure_tuple, ensure_tuple_rep, optional_import, require_pkg if TYPE_CHECKING: import itk import nibabel as nib + import nrrd + import pydicom from nibabel.nifti1 import Nifti1Image from PIL import Image as PILImage - has_itk = has_nib = has_pil = True + has_nrrd = has_itk = has_nib = has_pil = has_pydicom = True else: itk, has_itk = optional_import("itk", allow_namespace_pkg=True) nib, has_nib = optional_import("nibabel") Nifti1Image, _ = optional_import("nibabel.nifti1", name="Nifti1Image") PILImage, has_pil = optional_import("PIL.Image") + pydicom, has_pydicom = optional_import("pydicom") + nrrd, has_nrrd = optional_import("nrrd", allow_namespace_pkg=True) OpenSlide, _ = optional_import("openslide", name="OpenSlide") CuImage, _ = optional_import("cucim", name="CuImage") TiffFile, _ = optional_import("tifffile", name="TiffFile") -__all__ = ["ImageReader", "ITKReader", "NibabelReader", "NumpyReader", "PILReader", "WSIReader"] +__all__ = [ + "ImageReader", + "ITKReader", + "NibabelReader", + "NumpyReader", + "PILReader", + "PydicomReader", + "WSIReader", + "NrrdReader", +] class ImageReader(ABC): @@ -55,7 +76,7 @@ class ImageReader(ABC): img_data, meta_data = image_reader.get_data(img_obj) - The `read` call converts image filenames into image objects, - - The `get_data` call fetches the image data, as well as meta data. + - The `get_data` call fetches the image data, as well as metadata. - A reader should implement `verify_suffix` with the logic of checking the input filename by the filename extensions. @@ -91,9 +112,9 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs) -> Union[Seq @abstractmethod def get_data(self, img) -> Tuple[np.ndarray, Dict]: """ - Extract data array and meta data from loaded image and return them. + Extract data array and metadata from loaded image and return them. This function must return two objects, the first is a numpy array of image data, - the second is a dictionary of meta data. + the second is a dictionary of metadata. Args: img: an image object loaded from an image file or a list of image objects. @@ -147,7 +168,7 @@ class ITKReader(ImageReader): Args: channel_dim: the channel dimension of the input image, default is None. - This is used to set original_channel_dim in the meta data, EnsureChannelFirstD reads this field. + This is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field. If None, `original_channel_dim` will be either `no_channel` or `-1`. - Nifti file is usually "channel last", so there is no need to specify this argument. @@ -248,11 +269,11 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs): def get_data(self, img): """ - Extract data array and meta data from loaded image and return them. - This function returns two objects, first is numpy array of image data, second is dict of meta data. + Extract data array and metadata from loaded image and return them. + This function returns two objects, first is numpy array of image data, second is dict of metadata. It constructs `affine`, `original_affine`, and `spatial_shape` and stores them in meta dict. When loading a list of files, they are stacked together at a new dimension as the first dimension, - and the meta data of the first image is used to represent the output meta data. + and the metadata of the first image is used to represent the output metadata. Args: img: an ITK image object loaded from an image file or a list of ITK image objects. @@ -278,7 +299,7 @@ def get_data(self, img): def _get_meta_dict(self, img) -> Dict: """ - Get all the meta data of the image and convert to dict type. + Get all the metadata of the image and convert to dict type. Args: img: an ITK image object loaded from an image file. @@ -351,6 +372,312 @@ def _get_array_data(self, img): return np_img if self.reverse_indexing else np.moveaxis(np_img.T, 0, -1) +@require_pkg(pkg_name="pydicom") +class PydicomReader(ImageReader): + """ + Load medical images based on Pydicom library. + All the supported image formats can be found at: + https://dicom.nema.org/medical/dicom/current/output/chtml/part10/chapter_7.html + + This class refers to: + https://nipy.org/nibabel/dicom/dicom_orientation.html#dicom-affine-formula + https://github.com/pydicom/contrib-pydicom/blob/master/input-output/pydicom_series.py + + Args: + channel_dim: the channel dimension of the input image, default is None. + This is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field. + If None, `original_channel_dim` will be either `no_channel` or `-1`. + affine_lps_to_ras: whether to convert the affine matrix from "LPS" to "RAS". Defaults to ``True``. + Set to ``True`` to be consistent with ``NibabelReader``, + otherwise the affine matrix remains in the Dicom convention. + swap_ij: whether to swap the first two spatial axes. Default to ``True``, so that the outputs + are consistent with the other readers. + kwargs: additional args for `pydicom.dcmread` API. more details about available args: + https://pydicom.github.io/pydicom/stable/reference/generated/pydicom.filereader.dcmread.html#pydicom.filereader.dcmread + If the `get_data` function will be called + (for example, when using this reader with `monai.transforms.LoadImage`), please ensure that the argument + `stop_before_pixels` is `True`, and `specific_tags` covers all necessary tags, such as `PixelSpacing`, + `ImagePositionPatient`, `ImageOrientationPatient` and all `pixel_array` related tags. + """ + + def __init__( + self, channel_dim: Optional[int] = None, affine_lps_to_ras: bool = True, swap_ij: bool = True, **kwargs + ): + super().__init__() + self.kwargs = kwargs + self.channel_dim = channel_dim + self.affine_lps_to_ras = affine_lps_to_ras + self.swap_ij = swap_ij + + def verify_suffix(self, filename: Union[Sequence[PathLike], PathLike]) -> bool: + """ + Verify whether the specified file or files format is supported by Pydicom reader. + + Args: + filename: file name or a list of file names to read. + if a list of files, verify all the suffixes. + + """ + return has_pydicom + + def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs): + """ + Read image data from specified file or files, it can read a list of images + and stack them together as multi-channel data in `get_data()`. + If passing directory path instead of file path, will treat it as DICOM images series and read. + + Args: + data: file name or a list of file names to read, + kwargs: additional args for `pydicom.dcmread` API, will override `self.kwargs` for existing keys. + + Returns: + If `data` represents a filename: return a pydicom dataset object. + If `data` represents a list of filenames or a directory: return a list of pydicom dataset object. + If `data` represents a list of directories: return a list of list of pydicom dataset object. + + """ + img_ = [] + + filenames: Sequence[PathLike] = ensure_tuple(data) + kwargs_ = self.kwargs.copy() + kwargs_.update(kwargs) + + self.has_series = False + + for name in filenames: + name = f"{name}" + if Path(name).is_dir(): + # read DICOM series + slices = [pydicom.dcmread(fp=slc, **kwargs_) for slc in glob.glob(os.path.join(name, "**"))] + img_.append(slices if len(slices) > 1 else slices[0]) + self.has_series = True + else: + ds = pydicom.dcmread(fp=name, **kwargs_) + img_.append(ds) + return img_ if len(filenames) > 1 else img_[0] + + def _combine_dicom_series(self, data): + """ + Combine dicom series (a list of pydicom dataset objects). Their data arrays will be stacked together at a new + dimension as the last dimension. + + The stack order depends on Instance Number. The metadata will be based on the + first slice's metadata, and some new items will be added: + + "spacing": the new spacing of the stacked slices. + "lastImagePositionPatient": `ImagePositionPatient` for the last slice, it will be used to achieve the affine + matrix. + "spatial_shape": the spatial shape of the stacked slices. + + Args: + data: a list of pydicom dataset objects. + Returns: + a tuple that consisted with data array and metadata. + """ + slices = [] + # for a dicom series + for slc_ds in data: + if hasattr(slc_ds, "InstanceNumber"): + slices.append(slc_ds) + else: + warnings.warn(f"slice: {slc_ds.filename} does not have InstanceNumber tag, skip it.") + slices = sorted(slices, key=lambda s: s.InstanceNumber) + + if len(slices) == 0: + raise ValueError("the input does not have valid slices.") + + first_slice = slices[0] + average_distance = 0.0 + first_array = self._get_array_data(first_slice) + shape = first_array.shape + spacing = getattr(first_slice, "PixelSpacing", (1.0, 1.0, 1.0)) + pos = getattr(first_slice, "ImagePositionPatient", (0.0, 0.0, 0.0))[2] + stack_array = [first_array] + for idx in range(1, len(slices)): + slc_array = self._get_array_data(slices[idx]) + slc_shape = slc_array.shape + slc_spacing = getattr(first_slice, "PixelSpacing", (1.0, 1.0, 1.0)) + slc_pos = getattr(first_slice, "ImagePositionPatient", (0.0, 0.0, float(idx)))[2] + if spacing != slc_spacing: + warnings.warn(f"the list contains slices that have different spacings {spacing} and {slc_spacing}.") + if shape != slc_shape: + warnings.warn(f"the list contains slices that have different shapes {shape} and {slc_shape}.") + average_distance += abs(pos - slc_pos) + pos = slc_pos + stack_array.append(slc_array) + + if len(slices) > 1: + average_distance /= len(slices) - 1 + spacing.append(average_distance) + stack_array = np.stack(stack_array, axis=-1) + stack_metadata = self._get_meta_dict(first_slice) + stack_metadata["spacing"] = np.asarray(spacing) + if hasattr(slices[-1], "ImagePositionPatient"): + stack_metadata["lastImagePositionPatient"] = np.asarray(slices[-1].ImagePositionPatient) + stack_metadata["spatial_shape"] = shape + (len(slices),) + else: + stack_array = stack_array[0] + stack_metadata = self._get_meta_dict(first_slice) + stack_metadata["spacing"] = np.asarray(spacing) + stack_metadata["spatial_shape"] = shape + + return stack_array, stack_metadata + + def get_data(self, data): + """ + Extract data array and metadata from loaded image and return them. + This function returns two objects, first is numpy array of image data, second is dict of metadata. + It constructs `affine`, `original_affine`, and `spatial_shape` and stores them in meta dict. + For dicom series within the input, all slices will be stacked first, + When loading a list of files (dicom file, or stacked dicom series), they are stacked together at a new + dimension as the first dimension, and the metadata of the first image is used to represent the output metadata. + + To use this function, all pydicom dataset objects should contain: `pixel_array`, `PixelSpacing`, + `ImagePositionPatient` and `ImageOrientationPatient`. + + Args: + data: a pydicom dataset object, or a list of pydicom dataset objects, or a list of list of + pydicom dataset objects. + + """ + + dicom_data = [] + # combine dicom series if exists + if self.has_series is True: + # a list, all objects within a list belong to one dicom series + if not isinstance(data[0], List): + dicom_data.append(self._combine_dicom_series(data)) + # a list of list, each inner list represents a dicom series + else: + for series in data: + dicom_data.append(self._combine_dicom_series(series)) + else: + # a single pydicom dataset object + if not isinstance(data, List): + data = [data] + for d in data: + data_array = self._get_array_data(d) + metadata = self._get_meta_dict(d) + metadata["spatial_shape"] = data_array.shape + metadata["spacing"] = np.asarray(metadata.get("00280030", {}).get("Value", (1.0, 1.0, 1.0))) + dicom_data.append((data_array, metadata)) + + img_array: List[np.ndarray] = [] + compatible_meta: Dict = {} + + for (data_array, metadata) in ensure_tuple(dicom_data): + img_array.append(np.ascontiguousarray(np.swapaxes(data_array, 0, 1) if self.swap_ij else data_array)) + affine = self._get_affine(metadata, self.affine_lps_to_ras) + if self.swap_ij: + affine = affine @ np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) + sp_size = list(metadata["spatial_shape"]) + sp_size[0], sp_size[1] = sp_size[1], sp_size[0] + metadata["spatial_shape"] = ensure_tuple(sp_size) + metadata["original_affine"] = affine + metadata["affine"] = affine.copy() + if self.channel_dim is None: # default to "no_channel" or -1 + metadata["original_channel_dim"] = ( + "no_channel" if len(data_array.shape) == len(metadata["spatial_shape"]) else -1 + ) + else: + metadata["original_channel_dim"] = self.channel_dim + metadata["spacing"] = affine_to_spacing(metadata["original_affine"], r=len(metadata["spatial_shape"])) + _copy_compatible_dict(metadata, compatible_meta) + + return _stack_images(img_array, compatible_meta), compatible_meta + + def _get_meta_dict(self, img) -> Dict: + """ + Get all the metadata of the image and convert to dict type. + + Args: + img: a Pydicom dataset object. + + """ + if not hasattr(img, "ImagePositionPatient"): + raise ValueError(f"dicom data: {img.filename} does not have ImagePositionPatient.") + if not hasattr(img, "ImageOrientationPatient"): + raise ValueError(f"dicom data: {img.filename} does not have ImageOrientationPatient.") + + meta_dict = img.to_json_dict() + # remove Pixel Data "7FE00008" or "7FE00009" or "7FE00010" + # remove Data Set Trailing Padding "FFFCFFFC" + for key in ["7FE00008", "7FE00009", "7FE00010", "FFFCFFFC"]: + if key in meta_dict.keys(): + meta_dict.pop(key) + + return meta_dict # type: ignore + + def _get_affine(self, metadata: Dict, lps_to_ras: bool = True): + """ + Get or construct the affine matrix of the image, it can be used to correct + spacing, orientation or execute spatial transforms. + + Args: + metadata: metadata with dict type. + lps_to_ras: whether to convert the affine matrix from "LPS" to "RAS". Defaults to True. + + """ + affine: np.ndarray = np.eye(4) + if not ("00200037" in metadata and "00200032" in metadata): + return affine + # "00200037" is the tag of `ImageOrientationPatient` + rx, ry, rz, cx, cy, cz = metadata["00200037"]["Value"] + # "00200032" is the tag of `ImagePositionPatient` + sx, sy, sz = metadata["00200032"]["Value"] + dr, dc = metadata.get("spacing", (1.0, 1.0))[:2] + affine[0, 0] = cx * dr + affine[0, 1] = rx * dc + affine[0, 3] = sx + affine[1, 0] = cy * dr + affine[1, 1] = ry * dc + affine[1, 3] = sy + affine[2, 0] = cz * dr + affine[2, 1] = rz * dc + affine[2, 2] = 0 + affine[2, 3] = sz + + # 3d + if "lastImagePositionPatient" in metadata: + t1n, t2n, t3n = metadata["lastImagePositionPatient"] + n = metadata["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 lps_to_ras: + affine = orientation_ras_lps(affine) + return affine + + def _get_array_data(self, img): + """ + Get the array data of the image. If `RescaleSlope` and `RescaleIntercept` are available, the raw array data + will be rescaled. The output data has the dtype np.float32 if the rescaling is applied. + + Args: + img: a Pydicom dataset object. + + """ + # process Dicom series + if not hasattr(img, "pixel_array"): + raise ValueError(f"dicom data: {img.filename} does not have pixel_array.") + data = img.pixel_array + + slope, offset = 1.0, 0.0 + rescale_flag = False + if hasattr(img, "RescaleSlope"): + slope = img.RescaleSlope + rescale_flag = True + if hasattr(img, "RescaleIntercept"): + offset = img.RescaleIntercept + rescale_flag = True + if rescale_flag: + data = data.astype(np.float32) * slope + offset + + return data + + @require_pkg(pkg_name="nibabel") class NibabelReader(ImageReader): """ @@ -360,7 +687,7 @@ class NibabelReader(ImageReader): as_closest_canonical: if True, load the image as closest to canonical axis format. squeeze_non_spatial_dims: if True, non-spatial singletons will be squeezed, e.g. (256,256,1,3) -> (256,256,3) channel_dim: the channel dimension of the input image, default is None. - this is used to set original_channel_dim in the meta data, EnsureChannelFirstD reads this field. + this is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field. if None, `original_channel_dim` will be either `no_channel` or `-1`. most Nifti files are usually "channel last", no need to specify this argument for them. dtype: dtype of the output data array when loading with Nibabel library. @@ -422,11 +749,11 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs): def get_data(self, img): """ - Extract data array and meta data from loaded image and return them. - This function returns two objects, first is numpy array of image data, second is dict of meta data. + Extract data array and metadata from loaded image and return them. + This function returns two objects, first is numpy array of image data, second is dict of metadata. It constructs `affine`, `original_affine`, and `spatial_shape` and stores them in meta dict. When loading a list of files, they are stacked together at a new dimension as the first dimension, - and the meta data of the first image is used to present the output meta data. + and the metadata of the first image is used to present the output metadata. Args: img: a Nibabel image object loaded from an image file or a list of Nibabel image objects. @@ -460,7 +787,7 @@ def get_data(self, img): def _get_meta_dict(self, img) -> Dict: """ - Get the all the meta data of the image and convert to dict type. + Get the all the metadata of the image and convert to dict type. Args: img: a Nibabel image object loaded from an image file. @@ -587,11 +914,11 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs): def get_data(self, img): """ - Extract data array and meta data from loaded image and return them. - This function returns two objects, first is numpy array of image data, second is dict of meta data. + Extract data array and metadata from loaded image and return them. + This function returns two objects, first is numpy array of image data, second is dict of metadata. It constructs `affine`, `original_affine`, and `spatial_shape` and stores them in meta dict. When loading a list of files, they are stacked together at a new dimension as the first dimension, - and the meta data of the first image is used to represent the output meta data. + and the metadata of the first image is used to represent the output metadata. Args: img: a Numpy array loaded from a file or a list of Numpy arrays. @@ -673,11 +1000,11 @@ def read(self, data: Union[Sequence[PathLike], PathLike, np.ndarray], **kwargs): def get_data(self, img): """ - Extract data array and meta data from loaded image and return them. - This function returns two objects, first is numpy array of image data, second is dict of meta data. + Extract data array and metadata from loaded image and return them. + This function returns two objects, first is numpy array of image data, second is dict of metadata. It computes `spatial_shape` and stores it in meta dict. When loading a list of files, they are stacked together at a new dimension as the first dimension, - and the meta data of the first image is used to represent the output meta data. + and the metadata of the first image is used to represent the output metadata. Note that it will swap axis 0 and 1 after loading the array because the `HW` definition in PIL is different from other common medical packages. @@ -700,7 +1027,7 @@ def get_data(self, img): def _get_meta_dict(self, img) -> Dict: """ - Get the all the meta data of the image and convert to dict type. + Get the all the metadata of the image and convert to dict type. Args: img: a PIL Image object loaded from an image file. @@ -809,7 +1136,7 @@ def get_data( Args: img: a WSIReader image object loaded from a file, or list of CuImage objects - location: (x_min, y_min) tuple giving the top left pixel in the level 0 reference frame, + location: (top, left) tuple giving the top left pixel in the level 0 reference frame, or list of tuples (default=(0, 0)) size: (height, width) tuple giving the region size, or list of tuples (default to full image size) This is the size of image at the given level (`level`) @@ -820,7 +1147,10 @@ def get_data( """ # Verify inputs if level is None: - level = self._check_level(img, level) + level = self.level + max_level = self._get_max_level(img) + if level > max_level: + raise ValueError(f"The maximum level of this image is {max_level} while level={level} is requested)!") # Extract a region or the entire image region = self._extract_region(img, location=location, size=size, level=level, dtype=dtype) @@ -844,21 +1174,19 @@ def get_data( return patches, metadata - def _check_level(self, img, level): - level = self.level + def _get_max_level(self, img_obj): + """ + Return the maximum number of levels in the whole slide image + Args: + img: the whole slide image object - level_count = 0 + """ if self.backend == "openslide": - level_count = img.level_count - elif self.backend == "cucim": - level_count = img.resolutions["level_count"] - elif self.backend == "tifffile": - level_count = len(img.pages) - - if level > level_count - 1: - raise ValueError(f"The maximum level of this image is {level_count - 1} while level={level} is requested)!") - - return level + return img_obj.level_count - 1 + if self.backend == "cucim": + return img_obj.resolutions["level_count"] - 1 + if self.backend == "tifffile": + return len(img_obj.pages) - 1 def _get_image_size(self, img, size, level, location): """ @@ -975,3 +1303,158 @@ def _extract_patches( idx += 1 return flat_patch_grid + + +@dataclass +class NrrdImage: + """Class to wrap nrrd image array and metadata header""" + + array: np.ndarray + header: dict + + +@require_pkg(pkg_name="nrrd") +class NrrdReader(ImageReader): + """ + Load NRRD format images based on pynrrd library. + + Args: + channel_dim: the channel dimension of the input image, default is None. + This is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field. + If None, `original_channel_dim` will be either `no_channel` or `0`. + NRRD files are usually "channel first". + dtype: dtype of the data array when loading image. + index_order: Specify whether the returned data array should be in C-order (‘C’) or Fortran-order (‘F’). + Numpy is usually in C-order, but default on the NRRD header is F + kwargs: additional args for `nrrd.read` API. more details about available args: + https://github.com/mhe/pynrrd/blob/master/nrrd/reader.py + + """ + + def __init__( + self, + channel_dim: Optional[int] = None, + dtype: Union[np.dtype, type, str, None] = np.float32, + index_order: str = "F", + **kwargs, + ): + self.channel_dim = channel_dim + self.dtype = dtype + self.index_order = index_order + self.kwargs = kwargs + + def verify_suffix(self, filename: Union[Sequence[PathLike], PathLike]) -> bool: + """ + Verify whether the specified `filename` is supported by pynrrd reader. + + Args: + filename: file name or a list of file names to read. + if a list of files, verify all the suffixes. + + """ + suffixes: Sequence[str] = ["nrrd", "seg.nrrd"] + return has_nrrd and is_supported_format(filename, suffixes) + + def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs) -> Union[Sequence[Any], Any]: + """ + Read image data from specified file or files. + Note that it returns a data object or a sequence of data objects. + + Args: + data: file name or a list of file names to read. + kwargs: additional args for actual `read` API of 3rd party libs. + + """ + img_: List = [] + filenames: Sequence[PathLike] = ensure_tuple(data) + kwargs_ = self.kwargs.copy() + kwargs_.update(kwargs) + for name in filenames: + nrrd_image = NrrdImage(*nrrd.read(name, index_order=self.index_order, *kwargs_)) + img_.append(nrrd_image) + return img_ if len(filenames) > 1 else img_[0] + + def get_data(self, img: Union[NrrdImage, List[NrrdImage]]) -> Tuple[np.ndarray, Dict]: + """ + Extract data array and metadata from loaded image and return them. + This function must return two objects, the first is a numpy array of image data, + the second is a dictionary of metadata. + + Args: + img: a `NrrdImage` loaded from an image file or a list of image objects. + + """ + img_array: List[np.ndarray] = [] + compatible_meta: Dict = {} + + for i in ensure_tuple(img): + data = i.array.astype(self.dtype) + img_array.append(data) + header = dict(i.header) + if self.index_order == "C": + header = self._convert_f_to_c_order(header) + header["original_affine"] = self._get_affine(i) + header = self._switch_lps_ras(header) + header["affine"] = header["original_affine"].copy() + header["spatial_shape"] = header["sizes"] + [header.pop(k) for k in ("sizes", "space origin", "space directions")] # rm duplicated data in header + + if self.channel_dim is None: # default to "no_channel" or -1 + header["original_channel_dim"] = "no_channel" if len(data.shape) == len(header["spatial_shape"]) else 0 + else: + header["original_channel_dim"] = self.channel_dim + _copy_compatible_dict(header, compatible_meta) + + return _stack_images(img_array, compatible_meta), compatible_meta + + def _get_affine(self, img: NrrdImage) -> np.ndarray: + """ + Get the affine matrix of the image, it can be used to correct + spacing, orientation or execute spatial transforms. + + Args: + img: A `NrrdImage` loaded from image file + + """ + direction = img.header["space directions"] + origin = img.header["space origin"] + + x, y = direction.shape + affine_diam = min(x, y) + 1 + affine: np.ndarray = np.eye(affine_diam) + affine[:x, :y] = direction + affine[: (affine_diam - 1), -1] = origin # len origin is always affine_diam - 1 + return affine + + def _switch_lps_ras(self, header: dict) -> dict: + """ + For compatibility with nibabel, switch from LPS to RAS. Adapt affine matrix and + `space` argument in header accordingly. If no information of space is given in the header, + LPS is assumed and thus converted to RAS. If information about space is given, + but is not LPS, the unchanged header is returned. + + Args: + header: The image metadata as dict + + """ + if "space" not in header or header["space"] == "left-posterior-superior": + header["space"] = "right-anterior-superior" + header["original_affine"] = orientation_ras_lps(header["original_affine"]) + return header + + def _convert_f_to_c_order(self, header: dict) -> dict: + """ + All header fields of a NRRD are specified in `F` (Fortran) order, even if the image was read as C-ordered array. + 1D arrays of header['space origin'] and header['sizes'] become inverted, e.g, [1,2,3] -> [3,2,1] + The 2D Array for header['space directions'] is transposed: [[1,0,0],[0,2,0],[0,0,3]] -> [[3,0,0],[0,2,0],[0,0,1]] + For more details refer to: https://pynrrd.readthedocs.io/en/latest/user-guide.html#index-ordering + + Args: + header: The image metadata as dict + + """ + + header["space directions"] = np.rot90(np.flip(header["space directions"], 0)) + header["space origin"] = header["space origin"][::-1] + header["sizes"] = header["sizes"][::-1] + return header diff --git a/monai/data/image_writer.py b/monai/data/image_writer.py index cf9ef90e8c..27cfef6db1 100644 --- a/monai/data/image_writer.py +++ b/monai/data/image_writer.py @@ -15,6 +15,7 @@ from monai.apps.utils import get_logger from monai.config import DtypeLike, NdarrayOrTensor, PathLike +from monai.data.meta_tensor import MetaTensor from monai.data.utils import affine_to_spacing, ensure_tuple, ensure_tuple_rep, orientation_ras_lps, to_affine_nd from monai.transforms.spatial.array import Resize, SpatialResample from monai.transforms.utils_pytorch_numpy_unification import ascontiguousarray, moveaxis @@ -24,6 +25,7 @@ InterpolateMode, OptionalImportError, convert_data_type, + convert_to_tensor, look_up_option, optional_import, require_pkg, @@ -204,8 +206,8 @@ def resample_if_needed( affine: Optional[NdarrayOrTensor] = None, target_affine: Optional[NdarrayOrTensor] = None, output_spatial_shape: Union[Sequence[int], int, None] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: DtypeLike = np.float64, ): @@ -258,11 +260,18 @@ def resample_if_needed( ``np.float64`` for best precision. If ``None``, use the data type of input data. The output data type of this method is always ``np.float32``. """ + orig_type = type(data_array) + data_array = convert_to_tensor(data_array, track_meta=True) + if affine is not None: + data_array.affine = convert_to_tensor(affine, track_meta=False) # type: ignore resampler = SpatialResample(mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype) - output_array, target_affine = resampler( - data_array[None], src_affine=affine, dst_affine=target_affine, spatial_size=output_spatial_shape - ) - return output_array[0], target_affine + output_array = resampler(data_array[None], dst_affine=target_affine, spatial_size=output_spatial_shape) + # convert back at the end + if isinstance(output_array, MetaTensor): + output_array.applied_operations = [] + data_array, *_ = convert_data_type(output_array, output_type=orig_type) # type: ignore + affine, *_ = convert_data_type(output_array.affine, output_type=orig_type) # type: ignore + return data_array[0], affine @classmethod def convert_to_channel_last( @@ -613,6 +622,8 @@ def create_backend_obj( if dtype is not None: data_array = data_array.astype(dtype, copy=False) affine = convert_data_type(affine, np.ndarray)[0] + if affine is None: + affine = np.eye(4) affine = to_affine_nd(r=3, affine=affine) return nib.nifti1.Nifti1Image( data_array, @@ -736,7 +747,7 @@ def resample_and_clip( cls, data_array: NdarrayOrTensor, output_spatial_shape: Optional[Sequence[int]] = None, - mode: Union[InterpolateMode, str] = InterpolateMode.BICUBIC, + mode: str = InterpolateMode.BICUBIC, ): """ Resample ``data_array`` to ``output_spatial_shape`` if needed. @@ -755,11 +766,11 @@ def resample_and_clip( _min, _max = np.min(data), np.max(data) if len(data.shape) == 3: data = np.moveaxis(data, -1, 0) # to channel first - data = xform(data) # type: ignore + data = convert_data_type(xform(data), np.ndarray)[0] # type: ignore data = np.moveaxis(data, 0, -1) else: # (H, W) data = np.expand_dims(data, 0) # make a channel - data = xform(data)[0] # type: ignore + data = convert_data_type(xform(data), np.ndarray)[0][0] # type: ignore if mode != InterpolateMode.NEAREST: data = np.clip(data, _min, _max) return data @@ -792,7 +803,8 @@ def create_backend_obj( data: np.ndarray = super().create_backend_obj(data_array) if scale: # scale the data to be in an integer range - data = np.clip(data, 0.0, 1.0) # type: ignore # png writer only can scale data in range [0, 1] + data = np.clip(data, 0.0, 1.0) # png writer only can scale data in range [0, 1] + if scale == np.iinfo(np.uint8).max: data = (scale * data).astype(np.uint8, copy=False) elif scale == np.iinfo(np.uint16).max: diff --git a/monai/data/iterable_dataset.py b/monai/data/iterable_dataset.py index f292bf1593..f1906a80fe 100644 --- a/monai/data/iterable_dataset.py +++ b/monai/data/iterable_dataset.py @@ -11,7 +11,6 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Union -import numpy as np from torch.utils.data import IterableDataset as _TorchIterableDataset from torch.utils.data import get_worker_info @@ -115,9 +114,6 @@ def _get_item(): def randomize(self, size: int) -> None: self._idx = self.R.randint(size) - def set_random_state(self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None): - raise NotImplementedError(f"`set_random_state` is not available in {self.__class__.__name__}.") - class CSVIterableDataset(IterableDataset): """ diff --git a/monai/data/meta_obj.py b/monai/data/meta_obj.py new file mode 100644 index 0000000000..7d2e99ff79 --- /dev/null +++ b/monai/data/meta_obj.py @@ -0,0 +1,228 @@ +# 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 itertools +import pprint +from copy import deepcopy +from typing import Any, Iterable + +from monai.utils.enums import TraceKeys + +_TRACK_META = True + +__all__ = ["get_track_meta", "set_track_meta", "MetaObj"] + + +def set_track_meta(val: bool) -> None: + """ + Boolean to set whether metadata is tracked. If `True`, metadata will be associated + its data by using subclasses of `MetaObj`. If `False`, then data will be returned + with empty metadata. + + If `set_track_meta` is `False`, then standard data objects will be returned (e.g., + `torch.Tensor` and `np.ndarray`) as opposed to MONAI's enhanced objects. + + By default, this is `True`, and most users will want to leave it this way. However, + if you are experiencing any problems regarding metadata, and aren't interested in + preserving metadata, then you can disable it. + """ + global _TRACK_META + _TRACK_META = val + + +def get_track_meta() -> bool: + """ + Return the boolean as to whether metadata is tracked. If `True`, metadata will be + associated its data by using subclasses of `MetaObj`. If `False`, then data will be + returned with empty metadata. + + If `set_track_meta` is `False`, then standard data objects will be returned (e.g., + `torch.Tensor` and `np.ndarray`) as opposed to MONAI's enhanced objects. + + By default, this is `True`, and most users will want to leave it this way. However, + if you are experiencing any problems regarding metadata, and aren't interested in + preserving metadata, then you can disable it. + """ + return _TRACK_META + + +class MetaObj: + """ + Abstract base class that stores data as well as any extra metadata. + + This allows for subclassing `torch.Tensor` and `np.ndarray` through multiple inheritance. + + Metadata is stored in the form of a dictionary. + + Behavior should be the same as extended class (e.g., `torch.Tensor` or `np.ndarray`) + aside from the extended meta functionality. + + Copying of information: + + * For `c = a + b`, then auxiliary data (e.g., metadata) will be copied from the + first instance of `MetaObj` if `a.is_batch` is False + (For batched data, the metdata will be shallow copied for efficiency purposes). + + """ + + def __init__(self): + self._meta: dict = MetaObj.get_default_meta() + self._applied_operations: list = MetaObj.get_default_applied_operations() + self._is_batch: bool = False + + @staticmethod + def flatten_meta_objs(*args: Iterable): + """ + Recursively flatten input and yield all instances of `MetaObj`. + This means that for both `torch.add(a, b)`, `torch.stack([a, b])` (and + their numpy equivalents), we return `[a, b]` if both `a` and `b` are of type + `MetaObj`. + + Args: + args: Iterables of inputs to be flattened. + Returns: + list of nested `MetaObj` from input. + """ + for a in itertools.chain(*args): + if isinstance(a, (list, tuple)): + yield from MetaObj.flatten_meta_objs(a) + elif isinstance(a, MetaObj): + yield a + + def _copy_attr(self, attributes: list[str], input_objs, defaults: list, deep_copy: bool) -> None: + """ + Copy attributes from the first in a list of `MetaObj`. In the case of + `torch.add(a, b)`, both `a` and `b` could be `MetaObj` or something else, so + check them all. Copy the first to `self`. + + We also perform a deep copy of the data if desired. + + Args: + attributes: a sequence of strings corresponding to attributes to be copied (e.g., `['meta']`). + input_objs: an iterable of `MetaObj` instances. We'll copy the attribute from the first one + that contains that particular attribute. + defaults: If none of `input_objs` have the attribute that we're + interested in, then use this default value/function (e.g., `lambda: {}`.) + the defaults must be the same length as `attributes`. + deep_copy: whether to deep copy the corresponding attribute. + + Returns: + Returns `None`, but `self` should be updated to have the copied attribute. + """ + found = [False] * len(attributes) + for i, (idx, a) in itertools.product(input_objs, enumerate(attributes)): + if not found[idx] and hasattr(i, a): + setattr(self, a, deepcopy(getattr(i, a)) if deep_copy else getattr(i, a)) + found[idx] = True + if all(found): + return + for a, f, d in zip(attributes, found, defaults): + if not f: + setattr(self, a, d() if callable(defaults) else d) + return + + def _copy_meta(self, input_objs, deep_copy=False) -> None: + """ + Copy metadata from an iterable of `MetaObj` instances. For a given attribute, we copy the + adjunct data from the first element in the list containing that attribute. + + Args: + input_objs: list of `MetaObj` to copy data from. + + """ + self._copy_attr( + ["meta", "applied_operations"], + input_objs, + [MetaObj.get_default_meta(), MetaObj.get_default_applied_operations()], + deep_copy, + ) + + @staticmethod + def get_default_meta() -> dict: + """Get the default meta. + + Returns: + default metadata. + """ + return {} + + @staticmethod + def get_default_applied_operations() -> list: + """Get the default applied operations. + + Returns: + default applied operations. + """ + return [] + + def __repr__(self) -> str: + """String representation of class.""" + out: str = "\nMetaData\n" + if self.meta is not None: + out += "".join(f"\t{k}: {v}\n" for k, v in self.meta.items()) + else: + out += "None" + + out += "\nApplied operations\n" + if self.applied_operations is not None: + out += pprint.pformat(self.applied_operations, indent=2, compact=True, width=120) + else: + out += "None" + + out += f"\nIs batch?: {self.is_batch}" + + return out + + @property + def meta(self) -> dict: + """Get the meta.""" + return self._meta if hasattr(self, "_meta") else MetaObj.get_default_meta() + + @meta.setter + def meta(self, d) -> None: + """Set the meta.""" + if d == TraceKeys.NONE: + self._meta = MetaObj.get_default_meta() + self._meta = d + + @property + def applied_operations(self) -> list[dict]: + """Get the applied operations.""" + if hasattr(self, "_applied_operations"): + return self._applied_operations + return MetaObj.get_default_applied_operations() + + @applied_operations.setter + def applied_operations(self, t) -> None: + """Set the applied operations.""" + if t == TraceKeys.NONE: + # received no operations when decollating a batch + self._applied_operations = MetaObj.get_default_applied_operations() + return + self._applied_operations = t + + def push_applied_operation(self, t: Any) -> None: + self._applied_operations.append(t) + + def pop_applied_operation(self) -> Any: + return self._applied_operations.pop() + + @property + def is_batch(self) -> bool: + """Return whether object is part of batch or not.""" + return self._is_batch if hasattr(self, "_is_batch") else False + + @is_batch.setter + def is_batch(self, val: bool) -> None: + """Set whether object is part of batch or not.""" + self._is_batch = val diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py new file mode 100644 index 0000000000..d0818bd25e --- /dev/null +++ b/monai/data/meta_tensor.py @@ -0,0 +1,396 @@ +# 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 copy import deepcopy +from typing import Any, Sequence + +import numpy as np +import torch + +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.utils import look_up_option +from monai.utils.enums import PostFix +from monai.utils.type_conversion import convert_data_type, convert_to_tensor + +__all__ = ["MetaTensor"] + + +class MetaTensor(MetaObj, torch.Tensor): + """ + Class that inherits from both `torch.Tensor` and `MetaObj`, adding support for metadata. + + Metadata is stored in the form of a dictionary. Nested, an affine matrix will be + stored. This should be in the form of `torch.Tensor`. + + Behavior should be the same as `torch.Tensor` aside from the extended + meta functionality. + + Copying of information: + + * For `c = a + b`, then auxiliary data (e.g., metadata) will be copied from the + first instance of `MetaTensor` if `a.is_batch` is False + (For batched data, the metdata will be shallow copied for efficiency purposes). + + Example: + .. code-block:: python + + import torch + from monai.data import MetaTensor + + t = torch.tensor([1,2,3]) + affine = torch.as_tensor([[2,0,0,0], + [0,2,0,0], + [0,0,2,0], + [0,0,0,1]], dtype=torch.float64) + meta = {"some": "info"} + m = MetaTensor(t, affine=affine, meta=meta) + m2 = m + m + assert isinstance(m2, MetaTensor) + assert m2.meta["some"] == "info" + assert torch.all(m2.affine == affine) + + Notes: + - Requires pytorch 1.9 or newer for full compatibility. + - Older versions of pytorch (<=1.8), `torch.jit.trace(net, im)` may + not work if `im` is of type `MetaTensor`. This can be resolved with + `torch.jit.trace(net, im.as_tensor())`. + - For pytorch < 1.8, sharing `MetaTensor` instances across processes may not be supported. + - For pytorch < 1.9, next(iter(meta_tensor)) returns a torch.Tensor. + see: https://github.com/pytorch/pytorch/issues/54457 + - A warning will be raised if in the constructor `affine` is not `None` and + `meta` already contains the key `affine`. + - You can query whether the `MetaTensor` is a batch with the `is_batch` attribute. + - With a batch of data, `batch[0]` will return the 0th image + with the 0th metadata. When the batch dimension is non-singleton, e.g., + `batch[:, 0]`, `batch[..., -1]` and `batch[1:3]`, then all (or a subset in the + last example) of the metadata will be returned, and `is_batch` will return `True`. + - When creating a batch with this class, use `monai.data.DataLoader` as opposed + to `torch.utils.data.DataLoader`, as this will take care of collating the + metadata properly. + """ + + @staticmethod + def __new__( + cls, + x, + affine: torch.Tensor | None = None, + meta: dict | None = None, + applied_operations: list | None = None, + *args, + **kwargs, + ) -> MetaTensor: + _kwargs = {"device": kwargs.pop("device", None), "dtype": kwargs.pop("dtype", None)} if kwargs else {} + return torch.as_tensor(x, *args, **_kwargs).as_subclass(cls) # type: ignore + + def __init__( + self, + x, + affine: torch.Tensor | None = None, + meta: dict | None = None, + applied_operations: list | None = None, + *_args, + **_kwargs, + ) -> None: + """ + If `meta` is given, use it. Else, if `meta` exists in the input tensor, use it. + Else, use the default value. Similar for the affine, except this could come from + four places. + Priority: `affine`, `meta["affine"]`, `x.affine`, `get_default_affine`. + """ + super().__init__() + # set meta + if meta is not None: + self.meta = meta + elif isinstance(x, MetaObj): + self.meta = x.meta + # set the affine + if affine is not None: + if "affine" in self.meta: + warnings.warn("Setting affine, but the applied meta contains an affine. This will be overwritten.") + self.affine = affine + elif "affine" in self.meta: + # by using the setter function, we ensure it is converted to torch.Tensor if not already + self.affine = self.meta["affine"] + elif isinstance(x, MetaTensor): + self.affine = x.affine + else: + self.affine = self.get_default_affine() + # applied_operations + if applied_operations is not None: + self.applied_operations = applied_operations + elif isinstance(x, MetaTensor): + self.applied_operations = x.applied_operations + else: + self.applied_operations = MetaObj.get_default_applied_operations() + + # if we are creating a new MetaTensor, then deep copy attributes + if isinstance(x, torch.Tensor) and not isinstance(x, MetaTensor): + self.meta = deepcopy(self.meta) + self.applied_operations = deepcopy(self.applied_operations) + self.affine = self.affine.to(self.device) + + def _copy_attr(self, attributes: list[str], input_objs, defaults: list, deep_copy: bool) -> None: + super()._copy_attr(attributes, input_objs, defaults, deep_copy) + for a in attributes: + val = getattr(self, a) + if isinstance(val, torch.Tensor): + setattr(self, a, val.to(self.device)) + + @staticmethod + def update_meta(rets: Sequence, func, args, kwargs) -> Sequence: + """ + Update the metadata from the output of `MetaTensor.__torch_function__`. + + The output of `torch.Tensor.__torch_function__` could be a single object or a + sequence of them. Hence, in `MetaTensor.__torch_function__` we convert them to a + list of not already, and then we loop across each element, processing metadata + as necessary. For each element, if not of type `MetaTensor`, then nothing to do. + + Args: + rets: the output from `torch.Tensor.__torch_function__`, which has been + converted to a list in `MetaTensor.__torch_function__` if it wasn't + already a `Sequence`. + func: the torch function that was applied. Examples might be `torch.squeeze` + or `torch.Tensor.__add__`. We need this since the metadata need to be + treated differently if a batch of data is considered. For example, + slicing (`torch.Tensor.__getitem__`) the ith element of the 0th + dimension of a batch of data should return a ith tensor with the ith + metadata. + args: positional arguments that were passed to `func`. + kwargs: keyword arguments that were passed to `func`. + + Returns: + A sequence with the same number of elements as `rets`. For each element, if + the input type was not `MetaTensor`, then no modifications will have been + made. If global parameters have been set to false (e.g., + `not get_track_meta()`), then any `MetaTensor` will be converted to + `torch.Tensor`. Else, metadata will be propogated as necessary (see + :py:func:`MetaTensor._copy_meta`). + """ + out = [] + metas = None + is_batch = any(x.is_batch for x in MetaObj.flatten_meta_objs(args, kwargs.values()) if hasattr(x, "is_batch")) + for idx, ret in enumerate(rets): + # if not `MetaTensor`, nothing to do. + if not isinstance(ret, MetaTensor): + pass + # if not tracking, convert to `torch.Tensor`. + elif not get_track_meta(): + ret = ret.as_tensor() + # else, handle the `MetaTensor` metadata. + else: + meta_args = MetaObj.flatten_meta_objs(args, kwargs.values()) + ret._copy_meta(meta_args, deep_copy=not is_batch) + ret.is_batch = is_batch + # the following is not implemented but the network arch may run into this case: + # if func == torch.cat and any(m.is_batch if hasattr(m, "is_batch") else False for m in meta_args): + # raise NotImplementedError("torch.cat is not implemented for batch of MetaTensors.") + + # If we have a batch of data, then we need to be careful if a slice of + # the data is returned. Depending on how the data are indexed, we return + # some or all of the metadata, and the return object may or may not be a + # batch of data (e.g., `batch[:,-1]` versus `batch[0]`). + if is_batch: + # if indexing e.g., `batch[0]` + if func == torch.Tensor.__getitem__: + batch_idx = args[1] + if isinstance(batch_idx, Sequence): + batch_idx = batch_idx[0] + # if using e.g., `batch[:, -1]` or `batch[..., -1]`, then the + # first element will be `slice(None, None, None)` and `Ellipsis`, + # respectively. Don't need to do anything with the metadata. + if batch_idx not in (slice(None, None, None), Ellipsis, None): + # only decollate metadata once + if metas is None: + metas = decollate_batch(ret.meta) + meta = metas[batch_idx] + # if using e.g., `batch[0:2]`, then `is_batch` should still be + # `True`. Also re-collate the remaining elements. + if isinstance(meta, list): + ret.meta = list_data_collate(meta) + # if using e.g., `batch[0]` or `batch[0, 1]`, then return single + # element from batch, and set `is_batch` to `False`. + else: + ret.meta = meta + ret.is_batch = False + # `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: + if len(args) > 1: + dim = args[1] + elif "dim" in kwargs: + dim = kwargs["dim"] + else: + dim = 0 + if dim == 0: + if metas is None: + metas = decollate_batch(ret.meta) + ret.meta = metas[idx] + ret.is_batch = False + + ret.affine = ret.affine.to(ret.device) + out.append(ret) + # if the input was a tuple, then return it as a tuple + return tuple(out) if isinstance(rets, tuple) else out + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None) -> Any: + """Wraps all torch functions.""" + if kwargs is None: + kwargs = {} + ret = super().__torch_function__(func, types, args, kwargs) + # if `out` has been used as argument, metadata is not copied, nothing to do. + # if "out" in kwargs: + # return ret + # we might have 1 or multiple outputs. Might be MetaTensor, might be something + # else (e.g., `__repr__` returns a string). + # Convert to list (if necessary), process, and at end remove list if one was added. + if ( + hasattr(torch, "return_types") + and hasattr(func, "__name__") + and hasattr(torch.return_types, func.__name__) + and isinstance(getattr(torch.return_types, func.__name__), type) + and isinstance(ret, getattr(torch.return_types, func.__name__)) + ): + # for torch.max(torch.tensor(1.0), dim=0), the return type is named-tuple like + out_items = MetaTensor.update_meta(ret, func, args, kwargs) + for idx in range(ret.n_fields): + ret[idx].meta = out_items[idx].meta + ret[idx].applied_operations = out_items[idx].applied_operations + return ret + if isinstance(ret, (str, bytes)) or not isinstance(ret, Sequence): + ret = [ret] + unpack = True + else: + unpack = False + ret = MetaTensor.update_meta(ret, func, args, kwargs) + return ret[0] if unpack else ret + + def get_default_affine(self, dtype=torch.float64) -> torch.Tensor: + return torch.eye(4, device=self.device, dtype=dtype) + + def as_tensor(self) -> torch.Tensor: + """ + Return the `MetaTensor` as a `torch.Tensor`. + It is OS dependent as to whether this will be a deep copy or not. + """ + return self.as_subclass(torch.Tensor) # type: ignore + + def as_dict(self, key: str) -> dict: + """ + Get the object as a dictionary for backwards compatibility. + This method makes a copy of the objects. + + Args: + key: Base key to store main data. The key for the metadata will be + determined using `PostFix.meta`. + + Return: + A dictionary consisting of two keys, the main data (stored under `key`) and + the metadata. + """ + return { + key: self.as_tensor().clone().detach(), + PostFix.meta(key): deepcopy(self.meta), + PostFix.transforms(key): deepcopy(self.applied_operations), + } + + def astype(self, dtype, device=None, *unused_args, **unused_kwargs): + """ + Cast to ``dtype``, sharing data whenever possible. + + Args: + dtype: dtypes such as np.float32, torch.float, "np.float32", float. + device: the device if `dtype` is a torch data type. + unused_args: additional args (currently unused). + unused_kwargs: additional kwargs (currently unused). + + Returns: + data array instance + """ + if isinstance(dtype, str): + mod_str, *dtype = dtype.split(".", 1) + dtype = mod_str if not dtype else dtype[0] + else: + mod_str = getattr(dtype, "__module__", "torch") + mod_str = look_up_option(mod_str, {"torch", "numpy", "np"}, default="numpy") + if mod_str == "torch": + out_type = torch.Tensor + elif mod_str in ("numpy", "np"): + out_type = np.ndarray + else: + out_type = None + return convert_data_type(self, output_type=out_type, device=device, dtype=dtype, wrap_sequence=True)[0] + + @property + def affine(self) -> torch.Tensor: + """Get the affine.""" + return self.meta.get("affine", self.get_default_affine()) + + @affine.setter + def affine(self, d: NdarrayTensor) -> None: + """Set the affine.""" + self.meta["affine"] = torch.as_tensor(d, device=self.device) + + @property + def pixdim(self): + """Get the spacing""" + return affine_to_spacing(self.affine) + + def new_empty(self, size, dtype=None, device=None, requires_grad=False): + """ + must be defined for deepcopy to work + + See: + - https://pytorch.org/docs/stable/generated/torch.Tensor.new_empty.html#torch-tensor-new-empty + """ + return type(self)( + self.as_tensor().new_empty(size=size, dtype=dtype, device=device, requires_grad=requires_grad) + ) + + @staticmethod + def ensure_torch_and_prune_meta(im: NdarrayTensor, meta: dict, simple_keys: bool = False): + """ + Convert the image to `torch.Tensor`. If `affine` is in the `meta` dictionary, + convert that to `torch.Tensor`, too. Remove any superfluous metadata. + + Args: + im: Input image (`np.ndarray` or `torch.Tensor`) + meta: Metadata dictionary. + + Returns: + By default, a `MetaTensor` is returned. + However, if `get_track_meta()` is `False`, a `torch.Tensor` is returned. + """ + img = convert_to_tensor(im) # potentially ascontiguousarray + + # if not tracking metadata, return `torch.Tensor` + if not get_track_meta() or meta is None: + return img + + # remove any superfluous metadata. + if simple_keys: + # ensure affine is of type `torch.Tensor` + if "affine" in meta: + meta["affine"] = convert_to_tensor(meta["affine"]) # bc-breaking + remove_extra_metadata(meta) # bc-breaking + + # return the `MetaTensor` + return MetaTensor(img, meta=meta) + + def __repr__(self, *, tensor_contents=None): + return self.as_tensor().__repr__() + super().__repr__() diff --git a/monai/data/nifti_saver.py b/monai/data/nifti_saver.py index 3fdc0aa3e8..ddc5e10f63 100644 --- a/monai/data/nifti_saver.py +++ b/monai/data/nifti_saver.py @@ -29,8 +29,8 @@ class NiftiSaver: Typically, the data can be segmentation predictions, call `save` for single data or call `save_batch` to save a batch of data together. The name of saved file will be `{input_image_name}_{output_postfix}{output_ext}`, - where the input image name is extracted from the provided meta data dictionary. - If no meta data provided, use index from 0 as the filename prefix. + where the input image name is extracted from the provided metadata dictionary. + If no metadata provided, use index from 0 as the filename prefix. Note: image should include channel dimension: [B],C,H,W,[D]. @@ -45,8 +45,8 @@ def __init__( output_postfix: str = "seg", output_ext: str = ".nii.gz", resample: bool = True, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: DtypeLike = np.float64, output_dtype: DtypeLike = np.float32, @@ -99,8 +99,8 @@ def __init__( self.output_postfix = output_postfix self.output_ext = output_ext self.resample = resample - self.mode: GridSampleMode = GridSampleMode(mode) - self.padding_mode: GridSamplePadMode = GridSamplePadMode(padding_mode) + self.mode: str = GridSampleMode(mode) + self.padding_mode: str = GridSamplePadMode(padding_mode) self.align_corners = align_corners self.dtype = dtype self.output_dtype = output_dtype @@ -129,7 +129,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] Args: data: target data content that to be saved as a NIfTI format file. Assuming the data shape starts with a channel dimension and followed by spatial dimensions. - meta_data: the meta data information corresponding to the data. + meta_data: the metadata information corresponding to the data. See Also :py:meth:`monai.data.nifti_writer.write_nifti` diff --git a/monai/data/nifti_writer.py b/monai/data/nifti_writer.py index 8a6172955f..234f5b0a22 100644 --- a/monai/data/nifti_writer.py +++ b/monai/data/nifti_writer.py @@ -33,8 +33,8 @@ def write_nifti( target_affine: Optional[np.ndarray] = None, resample: bool = True, output_spatial_shape: Union[Sequence[int], np.ndarray, None] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: DtypeLike = np.float64, output_dtype: DtypeLike = np.float32, diff --git a/monai/data/png_saver.py b/monai/data/png_saver.py index 9a1ade0efa..5b6e3b5a30 100644 --- a/monai/data/png_saver.py +++ b/monai/data/png_saver.py @@ -28,8 +28,8 @@ class PNGSaver: Typically, the data can be segmentation predictions, call `save` for single data or call `save_batch` to save a batch of data together. The name of saved file will be `{input_image_name}_{output_postfix}{output_ext}`, - where the input image name is extracted from the provided meta data dictionary. - If no meta data provided, use index from 0 as the filename prefix. + where the input image name is extracted from the provided metadata dictionary. + If no metadata provided, use index from 0 as the filename prefix. .. deprecated:: 0.8 Use :py:class:`monai.transforms.SaveImage` instead. @@ -42,7 +42,7 @@ def __init__( output_postfix: str = "seg", output_ext: str = ".png", resample: bool = True, - mode: Union[InterpolateMode, str] = InterpolateMode.NEAREST, + mode: str = InterpolateMode.NEAREST, scale: Optional[int] = None, data_root_dir: PathLike = "", separate_folder: bool = True, @@ -103,7 +103,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] Assuming the data shape are spatial dimensions. Shape of the spatial dimensions (C,H,W). C should be 1, 3 or 4 - meta_data: the meta data information corresponding to the data. + meta_data: the metadata information corresponding to the data. Raises: ValueError: When ``data`` channels is not one of [1, 3, 4]. diff --git a/monai/data/png_writer.py b/monai/data/png_writer.py index 5d05536923..8c49944843 100644 --- a/monai/data/png_writer.py +++ b/monai/data/png_writer.py @@ -9,12 +9,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Sequence, Union +from typing import Optional, Sequence import numpy as np from monai.transforms.spatial.array import Resize -from monai.utils import InterpolateMode, deprecated, ensure_tuple_rep, look_up_option, optional_import +from monai.utils import ( + InterpolateMode, + convert_data_type, + deprecated, + ensure_tuple_rep, + look_up_option, + optional_import, +) Image, _ = optional_import("PIL", name="Image") @@ -24,7 +31,7 @@ def write_png( data: np.ndarray, file_name: str, output_spatial_shape: Optional[Sequence[int]] = None, - mode: Union[InterpolateMode, str] = InterpolateMode.BICUBIC, + mode: str = InterpolateMode.BICUBIC, scale: Optional[int] = None, ) -> None: """ @@ -38,7 +45,7 @@ def write_png( data: input data to write to file. file_name: expected file name that saved on disk. output_spatial_shape: spatial shape of the output image. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"bicubic"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html scale: {``255``, ``65535``} postprocess data by clipping to [0, 1] and scaling to @@ -74,9 +81,9 @@ def write_png( if scale is not None: data = np.clip(data, 0.0, 1.0) # png writer only can scale data in range [0, 1] if scale == np.iinfo(np.uint8).max: - data = (scale * data).astype(np.uint8, copy=False) + data = convert_data_type((scale * data), np.ndarray, dtype=np.uint8)[0] elif scale == np.iinfo(np.uint16).max: - data = (scale * data).astype(np.uint16, copy=False) + data = convert_data_type((scale * data), np.ndarray, dtype=np.uint16)[0] else: raise ValueError(f"Unsupported scale: {scale}, available options are [255, 65535]") diff --git a/monai/data/test_time_augmentation.py b/monai/data/test_time_augmentation.py index 0b97c9febf..3ce1fec2d4 100644 --- a/monai/data/test_time_augmentation.py +++ b/monai/data/test_time_augmentation.py @@ -72,11 +72,11 @@ class TestTimeAugmentation: image_key: key used to extract image from input dictionary. orig_key: the key of the original input data in the dict. will get the applied transform information for this input data, then invert them for the expected data with `image_key`. - orig_meta_keys: the key of the meta data of original input data, will get the `affine`, `data_shape`, etc. - the meta data is a dictionary object which contains: filename, original_shape, etc. + orig_meta_keys: the key of the metadata of original input data, will get the `affine`, `data_shape`, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`. - meta_key_postfix: use `key_{postfix}` to fetch the meta data according to the key data, - default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: use `key_{postfix}` to fetch the metadata according to the key data, + default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. this arg only works when `meta_keys=None`. diff --git a/monai/data/thread_buffer.py b/monai/data/thread_buffer.py index e21af69813..b882c2533c 100644 --- a/monai/data/thread_buffer.py +++ b/monai/data/thread_buffer.py @@ -10,9 +10,12 @@ # limitations under the License. +from multiprocessing.context import SpawnContext from queue import Empty, Full, Queue from threading import Thread +import torch + from monai.data import DataLoader, Dataset @@ -77,6 +80,60 @@ def __iter__(self): self.stop() # ensure thread completion +def buffer_iterator(src, buffer_size: int = 1, timeout: float = 0.01, repeats: int = 1): + """ + Create a ThreadBuffer object using the `src`, `buffer_size`, and `timeout` parameters given for the constructor + aguments of the same names, and yield each generated object `repeats` number of times successively. + + Args: + src: Source data iterable + buffer_size: Number of items to buffer from the source + timeout: Time to wait for an item from the buffer, or to wait while the buffer is full when adding items + repeats: Number of repeat generations to perform which is asynchronous from the generation of the next value + + Returns: + Generator yield (repeated) values from `src` asynchronously + """ + buffer = ThreadBuffer(src=src, buffer_size=buffer_size, timeout=timeout) + + for batch in buffer: + for _ in range(repeats): + yield batch + + +class _ProcessThread(Thread): + """Shim class to make a thread look like a process to the DataLoader class.""" + + @property + def pid(self): + return id(self) + + def run(self): + try: + super().run() + finally: + torch.utils.data._utils.worker._worker_info = None # clean up global data used for processes + + +class _ProcessQueue(Queue): + """Shim class to make a thread queue look like a process queue to the DataLoader class.""" + + def close(self): + pass + + def cancel_join_thread(self): + pass + + +class _ProcessThreadContext(SpawnContext): + _name = "processthread" + + # threads will be created which looks like processes + Process = _ProcessThread # type: ignore + # thread queue used in place of process queue to avoid some weird cleanup errors + Queue = _ProcessQueue # type: ignore + + class ThreadDataLoader(DataLoader): """ Subclass of `DataLoader` using a `ThreadBuffer` object to implement `__iter__` method asynchronously. This will @@ -87,7 +144,8 @@ class ThreadDataLoader(DataLoader): value the generated batch is yielded that many times while underlying dataset asynchronously generates the next. Typically not all relevant information is learned from a batch in a single iteration so training multiple times on the same batch will still produce good training with minimal short-term overfitting while allowing a slow batch - generation process more time to produce a result. + generation process more time to produce a result. This duplication is done by simply yielding the same object many + times and not by regenerating the data. Another typical usage is to accelerate light-weight preprocessing (usually cached all the deterministic transforms and no IO operations), because it leverages the separate thread to execute preprocessing to avoid unnecessary IPC @@ -95,6 +153,12 @@ class ThreadDataLoader(DataLoader): `ThreadDataLoader` can be useful for GPU transforms. For more details: https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md. + The `use_thread_workers` will cause workers to be created as threads rather than processes although everything else + in terms of how the class works is unchanged. This allows multiple workers to be used in Windows for example, or in + any other situation where thread semantics is desired. Please note that some MONAI components like several datasets + and random transforms are not thread-safe and can't work as expected with `thread workers`, need to check all the + preprocessing components carefully before enabling `use_thread_workers`. + See: * Fischetti et al. "Faster SGD training by minibatch persistency." ArXiv (2018) https://arxiv.org/abs/1806.07353 * Dami et al., "Faster Neural Network Training with Data Echoing" ArXiv (2020) https://arxiv.org/abs/1907.05550 @@ -106,21 +170,30 @@ class ThreadDataLoader(DataLoader): buffer_size: number of items to buffer from the data source. buffer_timeout: time to wait for an item from the buffer, or to wait while the buffer is full when adding items. repeats: number of times to yield the same batch. + use_thread_workers: if True and num_workers > 0 the workers are created as threads instead of processes kwargs: other arguments for `DataLoader` except for `dataset`. """ def __init__( - self, dataset: Dataset, buffer_size: int = 1, buffer_timeout: float = 0.01, repeats: int = 1, **kwargs + self, + dataset: Dataset, + buffer_size: int = 1, + buffer_timeout: float = 0.01, + repeats: int = 1, + use_thread_workers: bool = False, + **kwargs, ): + # if workers should be threads, create a new multiprocessing context with the process and queue types + # substituted with the shim types given above + if use_thread_workers and kwargs.get("num_workers", 0) > 0: + kwargs["multiprocessing_context"] = _ProcessThreadContext() + kwargs["persistent_workers"] = False + super().__init__(dataset, **kwargs) self.buffer_size = buffer_size self.buffer_timeout = buffer_timeout self.repeats = repeats def __iter__(self): - buffer = ThreadBuffer(src=super().__iter__(), buffer_size=self.buffer_size, timeout=self.buffer_timeout) - - for batch in buffer: - for _ in range(self.repeats): - yield batch + yield from buffer_iterator(super().__iter__(), self.buffer_size, self.buffer_timeout, self.repeats) diff --git a/monai/data/torchscript_utils.py b/monai/data/torchscript_utils.py index 61477e8ca9..ca46dd9dc4 100644 --- a/monai/data/torchscript_utils.py +++ b/monai/data/torchscript_utils.py @@ -18,7 +18,6 @@ from monai.config import get_config_values from monai.utils import JITMetadataKeys -from monai.utils.module import pytorch_after METADATA_FILENAME = "metadata.json" @@ -80,19 +79,10 @@ def save_net_with_metadata( json_data = json.dumps(metadict) - # Pytorch>1.6 can use dictionaries directly, otherwise need to use special map object - if pytorch_after(1, 7): - extra_files = {METADATA_FILENAME: json_data.encode()} + extra_files = {METADATA_FILENAME: json_data.encode()} - if more_extra_files is not None: - extra_files.update(more_extra_files) - else: - extra_files = torch._C.ExtraFilesMap() # type:ignore[attr-defined] - extra_files[METADATA_FILENAME] = json_data.encode() - - if more_extra_files is not None: - for k, v in more_extra_files.items(): - extra_files[k] = v + if more_extra_files is not None: + extra_files.update(more_extra_files) if isinstance(filename_prefix_or_stream, str): filename_no_ext, ext = os.path.splitext(filename_prefix_or_stream) @@ -123,16 +113,8 @@ def load_net_with_metadata( Returns: Triple containing loaded object, metadata dict, and extra files dict containing other file data if present """ - # Pytorch>1.6 can use dictionaries directly, otherwise need to use special map object - if pytorch_after(1, 7): - extra_files = {f: "" for f in more_extra_files} - extra_files[METADATA_FILENAME] = "" - else: - extra_files = torch._C.ExtraFilesMap() # type:ignore[attr-defined] - extra_files[METADATA_FILENAME] = "" - - for f in more_extra_files: - extra_files[f] = "" + extra_files = {f: "" for f in more_extra_files} + extra_files[METADATA_FILENAME] = "" jit_obj = torch.jit.load(filename_prefix_or_stream, map_location, extra_files) diff --git a/monai/data/utils.py b/monai/data/utils.py index 495daf15e2..45294cc66e 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -28,12 +28,14 @@ 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.networks.layers.simplelayers import GaussianFilter from monai.utils import ( MAX_SEED, BlendMode, Method, NumpyPadMode, + TraceKeys, convert_data_type, convert_to_dst_type, ensure_tuple, @@ -66,6 +68,7 @@ "get_valid_patch_size", "is_supported_format", "iter_patch", + "iter_patch_position", "iter_patch_slices", "json_hashing", "list_data_collate", @@ -84,6 +87,9 @@ "to_affine_nd", "worker_init_fn", "zoom_affine", + "remove_keys", + "remove_extra_metadata", + "get_extra_metadata_keys", ] # module to be used by `torch.save` @@ -119,33 +125,37 @@ def get_random_patch( def iter_patch_slices( - dims: Sequence[int], patch_size: Union[Sequence[int], int], start_pos: Sequence[int] = () + image_size: Sequence[int], + patch_size: Union[Sequence[int], int], + start_pos: Sequence[int] = (), + overlap: Union[Sequence[float], float] = 0.0, + padded: bool = True, ) -> Generator[Tuple[slice, ...], None, None]: """ - Yield successive tuples of slices defining patches of size `patch_size` from an array of dimensions `dims`. The - iteration starts from position `start_pos` in the array, or starting at the origin if this isn't provided. Each - patch is chosen in a contiguous grid using a first dimension as least significant ordering. + Yield successive tuples of slices defining patches of size `patch_size` from an array of dimensions `image_size`. + The iteration starts from position `start_pos` in the array, or starting at the origin if this isn't provided. Each + patch is chosen in a contiguous grid using a rwo-major ordering. Args: - dims: dimensions of array to iterate over + image_size: dimensions of array to iterate over patch_size: size of patches to generate slices for, 0 or None selects whole dimension start_pos: starting position in the array, default is 0 for each dimension + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + padded: if the image is padded so the patches can go beyond the borders. Defaults to False. Yields: Tuples of slice objects defining each patch """ - # ensure patchSize and startPos are the right length - ndim = len(dims) - patch_size_ = get_valid_patch_size(dims, patch_size) - start_pos = ensure_tuple_size(start_pos, ndim) - - # collect the ranges to step over each dimension - ranges = tuple(starmap(range, zip(start_pos, dims, patch_size_))) + # ensure patch_size has the right length + patch_size_ = get_valid_patch_size(image_size, patch_size) - # choose patches by applying product to the ranges - for position in product(*ranges[::-1]): # reverse ranges order to iterate in index order - yield tuple(slice(s, s + p) for s, p in zip(position[::-1], patch_size_)) + # create slices based on start position of each patch + for position in iter_patch_position( + image_size=image_size, patch_size=patch_size_, start_pos=start_pos, overlap=overlap, padded=padded + ): + yield tuple(slice(s, s + p) for s, p in zip(position, patch_size_)) def dense_patch_slices( @@ -188,12 +198,56 @@ def dense_patch_slices( return [tuple(slice(s, s + patch_size[d]) for d, s in enumerate(x)) for x in out] +def iter_patch_position( + image_size: Sequence[int], + patch_size: Union[Sequence[int], int], + start_pos: Sequence[int] = (), + overlap: Union[Sequence[float], float] = 0.0, + padded: bool = False, +): + """ + Yield successive tuples of upper left corner of patches of size `patch_size` from an array of dimensions `image_size`. + The iteration starts from position `start_pos` in the array, or starting at the origin if this isn't provided. Each + patch is chosen in a contiguous grid using a rwo-major ordering. + + Args: + image_size: dimensions of array to iterate over + patch_size: size of patches to generate slices for, 0 or None selects whole dimension + start_pos: starting position in the array, default is 0 for each dimension + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + padded: if the image is padded so the patches can go beyond the borders. Defaults to False. + + Yields: + Tuples of positions defining the upper left corner of each patch + """ + + # ensure patchSize and startPos are the right length + ndim = len(image_size) + patch_size_ = get_valid_patch_size(image_size, patch_size) + start_pos = ensure_tuple_size(start_pos, ndim) + overlap = ensure_tuple_rep(overlap, ndim) + + # calculate steps, which depends on the amount of overlap + steps = tuple(round(p * (1.0 - o)) for p, o in zip(patch_size_, overlap)) + + # calculate the last starting location (depending on the padding) + end_pos = image_size if padded else tuple(s - round(p) + 1 for s, p in zip(image_size, patch_size_)) + + # collect the ranges to step over each dimension + ranges = starmap(range, zip(start_pos, end_pos, steps)) + + # choose patches by applying product to the ranges + return product(*ranges) + + def iter_patch( arr: np.ndarray, patch_size: Union[Sequence[int], int] = 0, start_pos: Sequence[int] = (), + overlap: Union[Sequence[float], float] = 0.0, copy_back: bool = True, - mode: Union[NumpyPadMode, str] = NumpyPadMode.WRAP, + mode: Optional[str] = NumpyPadMode.WRAP, **pad_opts: Dict, ): """ @@ -205,11 +259,11 @@ def iter_patch( arr: array to iterate over patch_size: size of patches to generate slices for, 0 or None selects whole dimension start_pos: starting position in the array, default is 0 for each dimension + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. copy_back: if True data from the yielded patches is copied back to `arr` once the generator completes - mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, - ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - One of the listed string values or a user supplied function. Defaults to ``"wrap"``. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + mode: One of the listed string values in ``monai.utils.NumpyPadMode`` or ``monai.utils.PytorchPadMode``, + or a user supplied function. If None, no wrapping is performed. Defaults to ``"wrap"``. pad_opts: padding options, see `numpy.pad` Yields: @@ -229,19 +283,28 @@ def iter_patch( patch_size_ = get_valid_patch_size(arr.shape, patch_size) start_pos = ensure_tuple_size(start_pos, arr.ndim) + # set padded flag to false if pad mode is None + padded = True if mode else False # pad image by maximum values needed to ensure patches are taken from inside an image - arrpad = np.pad(arr, tuple((p, p) for p in patch_size_), look_up_option(mode, NumpyPadMode).value, **pad_opts) - - # choose a start position in the padded image - start_pos_padded = tuple(s + p for s, p in zip(start_pos, patch_size_)) - - # choose a size to iterate over which is smaller than the actual padded image to prevent producing - # patches which are only in the padded regions - iter_size = tuple(s + p for s, p in zip(arr.shape, patch_size_)) + if padded: + arrpad = np.pad(arr, tuple((p, p) for p in patch_size_), look_up_option(mode, NumpyPadMode).value, **pad_opts) + # choose a start position in the padded image + start_pos_padded = tuple(s + p for s, p in zip(start_pos, patch_size_)) + + # choose a size to iterate over which is smaller than the actual padded image to prevent producing + # patches which are only in the padded regions + iter_size = tuple(s + p for s, p in zip(arr.shape, patch_size_)) + else: + arrpad = arr + start_pos_padded = start_pos + iter_size = arr.shape - for slices in iter_patch_slices(iter_size, patch_size_, start_pos_padded): + for slices in iter_patch_slices(iter_size, patch_size_, start_pos_padded, overlap, padded=padded): # compensate original image padding - coords_no_pad = tuple((coord.start - p, coord.stop - p) for coord, p in zip(slices, patch_size_)) + if padded: + coords_no_pad = tuple((coord.start - p, coord.stop - p) for coord, p in zip(slices, patch_size_)) + else: + coords_no_pad = tuple((coord.start, coord.stop) for coord in slices) yield arrpad[slices], np.asarray(coords_no_pad) # data and coords (in numpy; works with torch loader) # copy back data from the padded image if required @@ -328,6 +391,24 @@ def dev_collate(batch, level: int = 1, logger_name: str = "dev_collate"): return +def collate_meta_tensor(batch): + """collate a sequence of meta tensor sequences/dictionaries into + a single batched metatensor or a dictionary of batched metatensor""" + if not isinstance(batch, Sequence): + raise NotImplementedError() + elem_0 = first(batch) + if isinstance(elem_0, MetaObj): + collated = default_collate(batch) + collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch]) + collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch] + collated.is_batch = True + return collated + if isinstance(elem_0, Mapping): + return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0} + # no more recursive search for MetaTensor + return default_collate(batch) + + def list_data_collate(batch: Sequence): """ Enhancement for PyTorch DataLoader default collate. @@ -346,9 +427,11 @@ def list_data_collate(batch: Sequence): ret = {} for k in elem: key = k - ret[key] = default_collate([d[key] for d in data]) - return ret - return default_collate(data) + data_for_batch = [d[key] for d in data] + ret[key] = collate_meta_tensor(data_for_batch) + else: + ret = collate_meta_tensor(data) + return ret except RuntimeError as re: re_str = str(re) if "equal size" in re_str: @@ -458,7 +541,9 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): """ if batch is None: return batch - if isinstance(batch, (float, int, str, bytes)): + if isinstance(batch, (float, int, str, bytes)) or ( + type(batch).__module__ == "numpy" and not isinstance(batch, Iterable) + ): return batch if isinstance(batch, torch.Tensor): if detach: @@ -466,6 +551,16 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): if batch.ndim == 0: return batch.item() if detach else batch out_list = torch.unbind(batch, dim=0) + # if of type MetaObj, decollate the metadata + if isinstance(batch, MetaObj): + for t, m in zip(out_list, decollate_batch(batch.meta)): + if isinstance(t, MetaObj): + t.meta = m + t.is_batch = False + for t, m in zip(out_list, batch.applied_operations): + if isinstance(t, MetaObj): + t.applied_operations = m + t.is_batch = False if out_list[0].ndim == 0 and detach: return [t.item() for t in out_list] return list(out_list) @@ -485,12 +580,7 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.") -def pad_list_data_collate( - batch: Sequence, - method: Union[Method, str] = Method.SYMMETRIC, - mode: Union[NumpyPadMode, str] = NumpyPadMode.CONSTANT, - **np_kwargs, -): +def pad_list_data_collate(batch: Sequence, method: str = Method.SYMMETRIC, mode: str = NumpyPadMode.CONSTANT, **kwargs): """ Function version of :py:class:`monai.transforms.croppad.batch.PadListDataCollate`. @@ -507,13 +597,13 @@ def pad_list_data_collate( batch: batch of data to pad-collate method: padding method (see :py:class:`monai.transforms.SpatialPad`) mode: padding mode (see :py:class:`monai.transforms.SpatialPad`) - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ from monai.transforms.croppad.batch import PadListDataCollate # needs to be here to avoid circular import - return PadListDataCollate(method=method, mode=mode, **np_kwargs)(batch) + return PadListDataCollate(method=method, mode=mode, **kwargs)(batch) def no_collation(x): @@ -566,6 +656,8 @@ def affine_to_spacing(affine: NdarrayTensor, r: int = 3, dtype=float, suppress_z Returns: an `r` dimensional vector of spacing. """ + if len(affine.shape) != 2 or affine.shape[0] != affine.shape[1]: + raise ValueError(f"affine must be a square matrix, got {affine.shape}.") _affine, *_ = convert_to_dst_type(affine[:r, :r], dst=affine, dtype=dtype) if isinstance(_affine, torch.Tensor): spacing = torch.sqrt(torch.sum(_affine * _affine, dim=0)) @@ -905,7 +997,7 @@ def compute_importance_map( mode = look_up_option(mode, BlendMode) device = torch.device(device) if mode == BlendMode.CONSTANT: - importance_map = torch.ones(patch_size, device=device).float() + importance_map = torch.ones(patch_size, device=device, dtype=torch.float) elif mode == BlendMode.GAUSSIAN: center_coords = [i // 2 for i in patch_size] sigma_scale = ensure_tuple_rep(sigma_scale, len(patch_size)) @@ -918,15 +1010,10 @@ def compute_importance_map( importance_map = importance_map.squeeze(0).squeeze(0) importance_map = importance_map / torch.max(importance_map) importance_map = importance_map.float() - - # importance_map cannot be 0, otherwise we may end up with nans! - min_non_zero = importance_map[importance_map != 0].min().item() - importance_map = torch.clamp(importance_map, min=min_non_zero) else: raise ValueError( f"Unsupported mode: {mode}, available options are [{BlendMode.CONSTANT}, {BlendMode.CONSTANT}]." ) - return importance_map @@ -1305,3 +1392,66 @@ def orientation_ras_lps(affine: NdarrayTensor) -> NdarrayTensor: if isinstance(affine, torch.Tensor): return torch.diag(torch.as_tensor(flip_diag).to(affine)) @ affine # type: ignore return np.diag(flip_diag).astype(affine.dtype) @ affine # type: ignore + + +def remove_keys(data: dict, keys: List[str]) -> None: + """ + Remove keys from a dictionary. Operates in-place so nothing is returned. + + Args: + data: dictionary to be modified. + keys: keys to be deleted from dictionary. + + Returns: + `None` + """ + for k in keys: + _ = data.pop(k, None) + + +def remove_extra_metadata(meta: dict) -> None: + """ + Remove extra metadata from the dictionary. Operates in-place so nothing is returned. + + Args: + meta: dictionary containing metadata to be modified. + + Returns: + `None` + """ + keys = get_extra_metadata_keys() + remove_keys(data=meta, keys=keys) + + +def get_extra_metadata_keys() -> List[str]: + """ + Get a list of unnecessary keys for metadata that can be removed. + + Returns: + List of keys to be removed. + """ + keys = [ + "srow_x", + "srow_y", + "srow_z", + "quatern_b", + "quatern_c", + "quatern_d", + "qoffset_x", + "qoffset_y", + "qoffset_z", + "dim", + "pixdim", + *[f"dim[{i}]" for i in range(8)], + *[f"pixdim[{i}]" for i in range(8)], + ] + + # TODO: it would be good to remove these, but they are currently being used in the + # codebase. + # keys += [ + # "original_affine", + # "spatial_shape", + # "spacing", + # ] + + return keys diff --git a/monai/data/wsi_datasets.py b/monai/data/wsi_datasets.py new file mode 100644 index 0000000000..37439c6e59 --- /dev/null +++ b/monai/data/wsi_datasets.py @@ -0,0 +1,402 @@ +# 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. + +import inspect +import os +from typing import Callable, Dict, Optional, Sequence, Tuple, Union + +import numpy as np + +from monai.data import Dataset +from monai.data.utils import iter_patch_position +from monai.data.wsi_reader import BaseWSIReader, WSIReader +from monai.transforms import ForegroundMask, Randomizable, apply_transform +from monai.utils import CommonKeys, ProbMapKeys, convert_to_dst_type, ensure_tuple_rep +from monai.utils.enums import WSIPatchKeys + +__all__ = ["PatchWSIDataset", "SlidingPatchWSIDataset", "MaskedPatchWSIDataset"] + + +class PatchWSIDataset(Dataset): + """ + This dataset extracts patches from whole slide images (without loading the whole image) + It also reads labels for each patch and provides each patch with its associated class labels. + + Args: + data: the list of input samples including image, location, and label (see the note below for more details). + size: the size of patch to be extracted from the whole slide image. + level: the level at which the patches to be extracted (default to 0). + transform: transforms to be executed on input data. + include_label: whether to load and include labels in the output + center_location: whether the input location information is the position of the center of the patch + additional_meta_keys: the list of keys for items to be copied to the output metadata from the input data + reader: the module to be used for loading whole slide imaging. If `reader` is + + - a string, it defines the backend of `monai.data.WSIReader`. Defaults to cuCIM. + - a class (inherited from `BaseWSIReader`), it is initialized and set as wsi_reader. + - an instance of a a class inherited from `BaseWSIReader`, it is set as the wsi_reader. + + kwargs: additional arguments to pass to `WSIReader` or provided whole slide reader class + + Note: + The input data has the following form as an example: + + .. code-block:: python + + [ + {"image": "path/to/image1.tiff", "patch_location": [200, 500], "label": 0}, + {"image": "path/to/image2.tiff", "patch_location": [100, 700], "patch_size": [20, 20], "patch_level": 2, "label": 1} + ] + + """ + + def __init__( + self, + data: Sequence, + patch_size: Optional[Union[int, Tuple[int, int]]] = None, + patch_level: Optional[int] = None, + transform: Optional[Callable] = None, + include_label: bool = True, + center_location: bool = True, + additional_meta_keys: Optional[Sequence[str]] = None, + reader="cuCIM", + **kwargs, + ): + super().__init__(data, transform) + + # Ensure patch size is a two dimensional tuple + if patch_size is None: + self.patch_size = None + else: + self.patch_size = ensure_tuple_rep(patch_size, 2) + + # Create a default level that override all levels if it is not None + self.patch_level = patch_level + # Set the default WSIReader's level to 0 if level is not provided + if patch_level is None: + patch_level = 0 + + # Setup the WSI reader + self.wsi_reader: Union[WSIReader, BaseWSIReader] + if isinstance(reader, str): + self.wsi_reader = WSIReader(backend=reader, level=patch_level, **kwargs) + elif inspect.isclass(reader) and issubclass(reader, BaseWSIReader): + self.wsi_reader = reader(level=patch_level, **kwargs) + elif isinstance(reader, BaseWSIReader): + self.wsi_reader = reader + else: + raise ValueError(f"Unsupported reader type: {reader}.") + self.backend = self.wsi_reader.backend + + self.include_label = include_label + self.center_location = center_location + self.additional_meta_keys = additional_meta_keys or [] + + # Initialized an empty whole slide image object dict + self.wsi_object_dict: Dict = {} + + def _get_wsi_object(self, sample: Dict): + image_path = sample[CommonKeys.IMAGE] + if image_path not in self.wsi_object_dict: + self.wsi_object_dict[image_path] = self.wsi_reader.read(image_path) + return self.wsi_object_dict[image_path] + + def _get_label(self, sample: Dict): + return np.array(sample[CommonKeys.LABEL], dtype=np.float32) + + def _get_location(self, sample: Dict): + if self.center_location: + size = self._get_size(sample) + return [sample[WSIPatchKeys.LOCATION][i] - size[i] // 2 for i in range(len(size))] + else: + return sample[WSIPatchKeys.LOCATION] + + def _get_level(self, sample: Dict): + if self.patch_level is None: + return sample.get(WSIPatchKeys.LEVEL, 0) + return self.patch_level + + def _get_size(self, sample: Dict): + if self.patch_size is None: + return ensure_tuple_rep(sample.get(WSIPatchKeys.SIZE), 2) + return self.patch_size + + def _get_data(self, sample: Dict): + # Don't store OpenSlide objects to avoid issues with OpenSlide internal cache + if self.backend == "openslide": + self.wsi_object_dict = {} + wsi_obj = self._get_wsi_object(sample) + location = self._get_location(sample) + level = self._get_level(sample) + size = self._get_size(sample) + return self.wsi_reader.get_data(wsi=wsi_obj, location=location, size=size, level=level) + + def _transform(self, index: int): + # Get a single entry of data + sample: Dict = self.data[index] + + # Extract patch image and associated metadata + image, metadata = self._get_data(sample) + output = {CommonKeys.IMAGE: image, CommonKeys.METADATA: metadata} + + # Include label in the output + if self.include_label: + output[CommonKeys.LABEL] = self._get_label(sample) + + for key in self.additional_meta_keys: + metadata[key] = sample[key] + + # Apply transforms and return it + return apply_transform(self.transform, output) if self.transform else output + + +class SlidingPatchWSIDataset(Randomizable, PatchWSIDataset): + """ + This dataset extracts patches from whole slide images (without loading the whole image) + It also reads labels for each patch and provides each patch with its associated class labels. + + Args: + data: the list of input samples including image, location, and label (see the note below for more details). + size: the size of patch to be extracted from the whole slide image. + level: the level at which the patches to be extracted (default to 0). + offset: the offset of image to extract patches (the starting position of the upper left patch). + offset_limits: if offset is set to "random", a tuple of integers defining the lower and upper limit of the + random offset for all dimensions, or a tuple of tuples that defines the limits for each dimension. + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + transform: transforms to be executed on input data. + reader: the module to be used for loading whole slide imaging. Defaults to cuCIM. If `reader` is + + - a string, it defines the backend of `monai.data.WSIReader`. + - a class (inherited from `BaseWSIReader`), it is initialized and set as wsi_reader, + - an instance of a a class inherited from `BaseWSIReader`, it is set as the wsi_reader. + + map_level: the resolution level at which the output map is created. + seed: random seed to randomly generate offsets. Defaults to 0. + kwargs: additional arguments to pass to `WSIReader` or provided whole slide reader class + + Note: + The input data has the following form as an example: + + .. code-block:: python + + [ + {"image": "path/to/image1.tiff"}, + {"image": "path/to/image2.tiff", "patch_size": [20, 20], "patch_level": 2} + ] + + """ + + def __init__( + self, + data: Sequence, + patch_size: Optional[Union[int, Tuple[int, int]]] = None, + patch_level: Optional[int] = None, + mask_level: int = 0, + overlap: Union[Tuple[float, float], float] = 0.0, + offset: Union[Tuple[int, int], int, str] = (0, 0), + offset_limits: Optional[Union[Tuple[Tuple[int, int], Tuple[int, int]], Tuple[int, int]]] = None, + transform: Optional[Callable] = None, + include_label: bool = False, + center_location: bool = False, + additional_meta_keys: Sequence[str] = (ProbMapKeys.LOCATION, ProbMapKeys.SIZE, ProbMapKeys.COUNT), + reader="cuCIM", + map_level: int = 0, + seed: int = 0, + **kwargs, + ): + super().__init__( + data=data, + patch_size=patch_size, + patch_level=patch_level, + transform=transform, + include_label=include_label, + center_location=center_location, + additional_meta_keys=additional_meta_keys, + reader=reader, + **kwargs, + ) + self.overlap = overlap + self.set_random_state(seed) + # Set the offset config + self.random_offset = False + if isinstance(offset, str): + if offset == "random": + self.random_offset = True + self.offset_limits: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] + if offset_limits is None: + self.offset_limits = None + elif isinstance(offset_limits, tuple): + if isinstance(offset_limits[0], int): + self.offset_limits = (offset_limits, offset_limits) + elif isinstance(offset_limits[0], tuple): + self.offset_limits = offset_limits + else: + raise ValueError( + "The offset limits should be either a tuple of integers or tuple of tuple of integers." + ) + else: + raise ValueError("The offset limits should be a tuple.") + else: + raise ValueError( + f'Invalid string for offset "{offset}". It should be either "random" as a string,' + "an integer, or a tuple of integers defining the offset." + ) + else: + self.offset = ensure_tuple_rep(offset, 2) + + self.mask_level = mask_level + # Create single sample for each patch (in a sliding window manner) + self.data = [] + self.image_data = data + for sample in self.image_data: + patch_samples = self._evaluate_patch_locations(sample) + self.data.extend(patch_samples) + + def _get_offset(self, sample): + if self.random_offset: + if self.offset_limits is None: + offset_limits = tuple((-s, s) for s in self._get_size(sample)) + else: + offset_limits = self.offset_limits + return tuple(self.R.randint(low, high) for low, high in offset_limits) + return self.offset + + def _evaluate_patch_locations(self, sample): + """Calculate the location for each patch in a sliding-window manner""" + patch_size = self._get_size(sample) + patch_level = self._get_level(sample) + wsi_obj = self._get_wsi_object(sample) + + # calculate the locations + wsi_size = self.wsi_reader.get_size(wsi_obj, 0) + mask_ratio = self.wsi_reader.get_downsample_ratio(wsi_obj, self.mask_level) + patch_ratio = self.wsi_reader.get_downsample_ratio(wsi_obj, patch_level) + patch_size_0 = np.array([p * patch_ratio for p in patch_size]) # patch size at level 0 + offset = self._get_offset(sample) + patch_locations = np.array( + list( + iter_patch_position( + image_size=wsi_size, patch_size=patch_size_0, start_pos=offset, overlap=self.overlap, padded=False + ) + ) + ) + # convert locations to mask_location + mask_locations = np.round((patch_locations + patch_size_0 // 2) / float(mask_ratio)) + + # fill out samples with location and metadata + sample[WSIPatchKeys.SIZE.value] = patch_size + sample[WSIPatchKeys.LEVEL.value] = patch_level + sample[ProbMapKeys.NAME.value] = os.path.basename(sample[CommonKeys.IMAGE]) + sample[ProbMapKeys.COUNT.value] = len(patch_locations) + sample[ProbMapKeys.SIZE.value] = np.array(self.wsi_reader.get_size(wsi_obj, self.mask_level)) + return [ + {**sample, WSIPatchKeys.LOCATION.value: np.array(loc), ProbMapKeys.LOCATION.value: mask_loc} + for loc, mask_loc in zip(patch_locations, mask_locations) + ] + + +class MaskedPatchWSIDataset(PatchWSIDataset): + """ + This dataset extracts patches from whole slide images at the locations where foreground mask + at a given level is non-zero. + + Args: + data: the list of input samples including image, location, and label (see the note below for more details). + size: the size of patch to be extracted from the whole slide image. + level: the level at which the patches to be extracted (default to 0). + mask_level: the resolution level at which the mask is created. + transform: transforms to be executed on input data. + include_label: whether to load and include labels in the output + center_location: whether the input location information is the position of the center of the patch + additional_meta_keys: the list of keys for items to be copied to the output metadata from the input data + reader: the module to be used for loading whole slide imaging. Defaults to cuCIM. If `reader` is + + - a string, it defines the backend of `monai.data.WSIReader`. + - a class (inherited from `BaseWSIReader`), it is initialized and set as wsi_reader, + - an instance of a a class inherited from `BaseWSIReader`, it is set as the wsi_reader. + + kwargs: additional arguments to pass to `WSIReader` or provided whole slide reader class + + Note: + The input data has the following form as an example: + + .. code-block:: python + + [ + {"image": "path/to/image1.tiff"}, + {"image": "path/to/image2.tiff", "patch_size": [20, 20], "patch_level": 2} + ] + + """ + + def __init__( + self, + data: Sequence, + patch_size: Optional[Union[int, Tuple[int, int]]] = None, + patch_level: Optional[int] = None, + mask_level: int = 7, + transform: Optional[Callable] = None, + include_label: bool = False, + center_location: bool = False, + additional_meta_keys: Sequence[str] = (ProbMapKeys.LOCATION, ProbMapKeys.NAME), + reader="cuCIM", + **kwargs, + ): + super().__init__( + data=data, + patch_size=patch_size, + patch_level=patch_level, + transform=transform, + include_label=include_label, + center_location=center_location, + additional_meta_keys=additional_meta_keys, + reader=reader, + **kwargs, + ) + + self.mask_level = mask_level + # Create single sample for each patch (in a sliding window manner) + self.data = [] + self.image_data = data + for sample in self.image_data: + patch_samples = self._evaluate_patch_locations(sample) + self.data.extend(patch_samples) + + def _evaluate_patch_locations(self, sample): + """Calculate the location for each patch based on the mask at different resolution level""" + patch_size = self._get_size(sample) + patch_level = self._get_level(sample) + wsi_obj = self._get_wsi_object(sample) + + # load the entire image at level=mask_level + wsi, _ = self.wsi_reader.get_data(wsi_obj, level=self.mask_level) + + # create the foreground tissue mask and get all indices for non-zero pixels + mask = np.squeeze(convert_to_dst_type(ForegroundMask(hsv_threshold={"S": "otsu"})(wsi), dst=wsi)[0]) + mask_locations = np.vstack(mask.nonzero()).T + + # convert mask locations to image locations at level=0 + mask_ratio = self.wsi_reader.get_downsample_ratio(wsi_obj, self.mask_level) + patch_ratio = self.wsi_reader.get_downsample_ratio(wsi_obj, patch_level) + patch_size_0 = np.array([p * patch_ratio for p in patch_size]) # patch size at level 0 + patch_locations = np.round((mask_locations + 0.5) * float(mask_ratio) - patch_size_0 // 2).astype(int) + + # fill out samples with location and metadata + sample[WSIPatchKeys.SIZE.value] = patch_size + sample[WSIPatchKeys.LEVEL.value] = patch_level + sample[ProbMapKeys.NAME.value] = os.path.basename(sample[CommonKeys.IMAGE]) + sample[ProbMapKeys.COUNT.value] = len(patch_locations) + sample[ProbMapKeys.SIZE.value] = mask.shape + return [ + {**sample, WSIPatchKeys.LOCATION.value: np.array(loc), ProbMapKeys.LOCATION.value: mask_loc} + for loc, mask_loc in zip(patch_locations, mask_locations) + ] diff --git a/monai/data/wsi_reader.py b/monai/data/wsi_reader.py new file mode 100644 index 0000000000..0bb2de987c --- /dev/null +++ b/monai/data/wsi_reader.py @@ -0,0 +1,581 @@ +# 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 abc import abstractmethod +from os.path import abspath +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import numpy as np + +from monai.config import DtypeLike, PathLike +from monai.data.image_reader import ImageReader, _stack_images +from monai.data.utils import is_supported_format +from monai.utils import WSIPatchKeys, ensure_tuple, optional_import, require_pkg + +CuImage, _ = optional_import("cucim", name="CuImage") +OpenSlide, _ = optional_import("openslide", name="OpenSlide") + +__all__ = ["BaseWSIReader", "WSIReader", "CuCIMWSIReader", "OpenSlideWSIReader"] + + +class BaseWSIReader(ImageReader): + """ + An abstract class that defines APIs to load patches from whole slide image files. + + Typical usage of a concrete implementation of this class is: + + .. code-block:: python + + image_reader = MyWSIReader() + wsi = image_reader.read(, **kwargs) + img_data, meta_data = image_reader.get_data(wsi) + + - The `read` call converts an image filename into whole slide image object, + - The `get_data` call fetches the image data, as well as metadata. + + The following methods needs to be implemented for any concrete implementation of this class: + + - `read` reads a whole slide image object from a given file + - `get_size` returns the size of the whole slide image of a given wsi object at a given level. + - `get_level_count` returns the number of levels in the whole slide image + - `get_patch` extracts and returns a patch image form the whole slide image + - `get_metadata` extracts and returns metadata for a whole slide image and a specific patch. + + + """ + + supported_suffixes: List[str] = [] + backend = "" + + def __init__(self, level: int, channel_dim: int = 0, **kwargs): + super().__init__() + self.level = level + self.channel_dim = channel_dim + self.kwargs = kwargs + self.metadata: Dict[Any, Any] = {} + + @abstractmethod + def get_size(self, wsi, level: int) -> Tuple[int, int]: + """ + Returns the size (height, width) of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + @abstractmethod + def get_level_count(self, wsi) -> int: + """ + Returns the number of levels in the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file + + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + @abstractmethod + def get_downsample_ratio(self, wsi, level: int) -> float: + """ + Returns the down-sampling ratio of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + @abstractmethod + def get_file_path(self, wsi) -> str: + """Return the file path for the WSI object""" + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + @abstractmethod + def get_patch( + self, wsi, location: Tuple[int, int], size: Tuple[int, int], level: int, dtype: DtypeLike, mode: str + ) -> np.ndarray: + """ + Extracts and returns a patch image form the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file or a lis of such objects + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + dtype: the data type of output image + mode: the output image mode, 'RGB' or 'RGBA' + + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") + + def get_metadata( + self, wsi, patch: np.ndarray, location: Tuple[int, int], size: Tuple[int, int], level: int + ) -> Dict: + """ + Returns metadata of the extracted patch from the whole slide image. + + Args: + wsi: the whole slide image object, from which the patch is loaded + patch: extracted patch from whole slide image + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + + """ + if self.channel_dim >= len(patch.shape) or self.channel_dim < -len(patch.shape): + ValueError(f"The desired channel_dim ({self.channel_dim}) is out of bound for image shape: {patch.shape}") + channel_dim: int = self.channel_dim + (len(patch.shape) if self.channel_dim < 0 else 0) + metadata: Dict = { + "backend": self.backend, + "original_channel_dim": channel_dim, + "spatial_shape": np.array(patch.shape[:channel_dim] + patch.shape[channel_dim + 1 :]), + "num_patches": 1, + WSIPatchKeys.PATH.value: self.get_file_path(wsi), + WSIPatchKeys.LOCATION.value: np.asarray(location), + WSIPatchKeys.SIZE.value: np.asarray(size), + WSIPatchKeys.LEVEL.value: level, + } + return metadata + + def get_data( + self, + wsi, + location: Tuple[int, int] = (0, 0), + size: Optional[Tuple[int, int]] = None, + level: Optional[int] = None, + dtype: DtypeLike = np.uint8, + mode: str = "RGB", + ) -> Tuple[np.ndarray, Dict]: + """ + Verifies inputs, extracts patches from WSI image and generates metadata, and return them. + + Args: + wsi: a whole slide image object loaded from a file or a list of such objects + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + dtype: the data type of output image + mode: the output image mode, 'RGB' or 'RGBA' + + Returns: + a tuples, where the first element is an image patch [CxHxW] or stack of patches, + and second element is a dictionary of metadata + """ + patch_list: List = [] + metadata_list: List = [] + # CuImage object is iterable, so ensure_tuple won't work on single object + if not isinstance(wsi, List): + wsi = [wsi] + for each_wsi in ensure_tuple(wsi): + # Verify magnification level + if level is None: + level = self.level + max_level = self.get_level_count(each_wsi) - 1 + if level > max_level: + raise ValueError(f"The maximum level of this image is {max_level} while level={level} is requested)!") + + # Verify location + if location is None: + location = (0, 0) + wsi_size = self.get_size(each_wsi, 0) + if location[0] > wsi_size[0] or location[1] > wsi_size[1]: + raise ValueError(f"Location is outside of the image: location={location}, image size={wsi_size}") + + # Verify size + if size is None: + if location != (0, 0): + raise ValueError("Patch size should be defined to extract patches.") + size = self.get_size(each_wsi, level) + else: + if size[0] <= 0 or size[1] <= 0: + raise ValueError(f"Patch size should be greater than zero, provided: patch size = {size}") + + # Extract a patch or the entire image + patch = self.get_patch(each_wsi, location=location, size=size, level=level, dtype=dtype, mode=mode) + + # check if the image has three dimensions (2D + color) + if patch.ndim != 3: + raise ValueError( + f"The image dimension should be 3 but has {patch.ndim}. " + "`WSIReader` is designed to work only with 2D images with color channel." + ) + # Check if there are four color channels for RGBA + if mode == "RGBA": + if patch.shape[self.channel_dim] != 4: + raise ValueError( + f"The image is expected to have four color channels in '{mode}' mode but has " + f"{patch.shape[self.channel_dim]}." + ) + # Check if there are three color channels for RGB + elif mode in "RGB" and patch.shape[self.channel_dim] != 3: + raise ValueError( + f"The image is expected to have three color channels in '{mode}' mode but has " + f"{patch.shape[self.channel_dim]}. " + ) + # Get patch-related metadata + metadata: dict = self.get_metadata(wsi=each_wsi, patch=patch, location=location, size=size, level=level) + # Create a list of patches and metadata + patch_list.append(patch) + metadata_list.append(metadata) + if len(wsi) > 1: + if len({m["original_channel_dim"] for m in metadata_list}) > 1: + raise ValueError("original_channel_dim is not consistent across wsi objects.") + if len({tuple(m["spatial_shape"]) for m in metadata_list}) > 1: + raise ValueError("spatial_shape is not consistent across wsi objects.") + for key in WSIPatchKeys: + metadata[key] = [m[key] for m in metadata_list] + return _stack_images(patch_list, metadata), metadata + + def verify_suffix(self, filename: Union[Sequence[PathLike], PathLike]) -> bool: + """ + Verify whether the specified file or files format is supported by WSI reader. + + The list of supported suffixes are read from `self.supported_suffixes`. + + Args: + filename: filename or a list of filenames to read. + + """ + return is_supported_format(filename, self.supported_suffixes) + + +class WSIReader(BaseWSIReader): + """ + Read whole slide images and extract patches using different backend libraries + + Args: + backend: the name of backend whole slide image reader library, the default is cuCIM. + level: the level at which patches are extracted. + channel_dim: the desired dimension for color channel. Default to 0 (channel first). + kwargs: additional arguments to be passed to the backend library + + """ + + def __init__(self, backend="cucim", level: int = 0, channel_dim: int = 0, **kwargs): + super().__init__(level, channel_dim, **kwargs) + self.backend = backend.lower() + self.reader: Union[CuCIMWSIReader, OpenSlideWSIReader] + if self.backend == "cucim": + self.reader = CuCIMWSIReader(level=level, channel_dim=channel_dim, **kwargs) + elif self.backend == "openslide": + self.reader = OpenSlideWSIReader(level=level, channel_dim=channel_dim, **kwargs) + else: + raise ValueError(f"The supported backends are cucim and openslide, '{self.backend}' was given.") + self.supported_suffixes = self.reader.supported_suffixes + + def get_level_count(self, wsi) -> int: + """ + Returns the number of levels in the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file + + """ + return self.reader.get_level_count(wsi) + + def get_size(self, wsi, level: int) -> Tuple[int, int]: + """ + Returns the size (height, width) of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return self.reader.get_size(wsi, level) + + def get_downsample_ratio(self, wsi, level: int) -> float: + """ + Returns the down-sampling ratio of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return self.reader.get_downsample_ratio(wsi, level) + + def get_file_path(self, wsi) -> str: + """Return the file path for the WSI object""" + return self.reader.get_file_path(wsi) + + def get_patch( + self, wsi, location: Tuple[int, int], size: Tuple[int, int], level: int, dtype: DtypeLike, mode: str + ) -> np.ndarray: + """ + Extracts and returns a patch image form the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file or a lis of such objects + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + dtype: the data type of output image + mode: the output image mode, 'RGB' or 'RGBA' + + """ + return self.reader.get_patch(wsi=wsi, location=location, size=size, level=level, dtype=dtype, mode=mode) + + def read(self, data: Union[Sequence[PathLike], PathLike, np.ndarray], **kwargs): + """ + Read whole slide image objects from given file or list of files. + + Args: + data: file name or a list of file names to read. + kwargs: additional args for the reader module (overrides `self.kwargs` for existing keys). + + Returns: + whole slide image object or list of such objects + + """ + return self.reader.read(data=data, **kwargs) + + +@require_pkg(pkg_name="cucim") +class CuCIMWSIReader(BaseWSIReader): + """ + Read whole slide images and extract patches using cuCIM library. + + Args: + level: the whole slide image level at which the image is extracted. (default=0) + This is overridden if the level argument is provided in `get_data`. + channel_dim: the desired dimension for color channel. Default to 0 (channel first). + kwargs: additional args for `cucim.CuImage` module: + https://github.com/rapidsai/cucim/blob/main/cpp/include/cucim/cuimage.h + + """ + + supported_suffixes = ["tif", "tiff", "svs"] + backend = "cucim" + + def __init__(self, level: int = 0, channel_dim: int = 0, **kwargs): + super().__init__(level, channel_dim, **kwargs) + + @staticmethod + def get_level_count(wsi) -> int: + """ + Returns the number of levels in the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file + + """ + return wsi.resolutions["level_count"] # type: ignore + + @staticmethod + def get_size(wsi, level: int) -> Tuple[int, int]: + """ + Returns the size (height, width) of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return (wsi.resolutions["level_dimensions"][level][1], wsi.resolutions["level_dimensions"][level][0]) + + @staticmethod + def get_downsample_ratio(wsi, level: int) -> float: + """ + Returns the down-sampling ratio of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return wsi.resolutions["level_downsamples"][level] # type: ignore + + def get_file_path(self, wsi) -> str: + """Return the file path for the WSI object""" + return str(abspath(wsi.path)) + + def read(self, data: Union[Sequence[PathLike], PathLike, np.ndarray], **kwargs): + """ + Read whole slide image objects from given file or list of files. + + Args: + data: file name or a list of file names to read. + kwargs: additional args that overrides `self.kwargs` for existing keys. + For more details look at https://github.com/rapidsai/cucim/blob/main/cpp/include/cucim/cuimage.h + + Returns: + whole slide image object or list of such objects + + """ + wsi_list: List = [] + + filenames: Sequence[PathLike] = ensure_tuple(data) + kwargs_ = self.kwargs.copy() + kwargs_.update(kwargs) + for filename in filenames: + wsi = CuImage(filename, **kwargs_) + wsi_list.append(wsi) + + return wsi_list if len(filenames) > 1 else wsi_list[0] + + def get_patch( + self, wsi, location: Tuple[int, int], size: Tuple[int, int], level: int, dtype: DtypeLike, mode: str + ) -> np.ndarray: + """ + Extracts and returns a patch image form the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file or a lis of such objects + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + dtype: the data type of output image + mode: the output image mode, 'RGB' or 'RGBA' + + """ + # Extract a patch or the entire image + # (reverse the order of location and size to become WxH for cuCIM) + patch: np.ndarray = wsi.read_region(location=location[::-1], size=size[::-1], level=level) + + # Convert to numpy + patch = np.asarray(patch, dtype=dtype) + + # Make the channel to desired dimensions + patch = np.moveaxis(patch, -1, self.channel_dim) + + # Check if the color channel is 3 (RGB) or 4 (RGBA) + if mode in "RGB": + if patch.shape[self.channel_dim] not in [3, 4]: + raise ValueError( + f"The image is expected to have three or four color channels in '{mode}' mode but has " + f"{patch.shape[self.channel_dim]}. " + ) + patch = patch[:3] + + return patch + + +@require_pkg(pkg_name="openslide") +class OpenSlideWSIReader(BaseWSIReader): + """ + Read whole slide images and extract patches using OpenSlide library. + + Args: + level: the whole slide image level at which the image is extracted. (default=0) + This is overridden if the level argument is provided in `get_data`. + channel_dim: the desired dimension for color channel. Default to 0 (channel first). + kwargs: additional args for `openslide.OpenSlide` module. + + """ + + supported_suffixes = ["tif", "tiff", "svs"] + backend = "openslide" + + def __init__(self, level: int = 0, channel_dim: int = 0, **kwargs): + super().__init__(level, channel_dim, **kwargs) + + @staticmethod + def get_level_count(wsi) -> int: + """ + Returns the number of levels in the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file + + """ + return wsi.level_count # type: ignore + + @staticmethod + def get_size(wsi, level: int) -> Tuple[int, int]: + """ + Returns the size (height, width) of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return (wsi.level_dimensions[level][1], wsi.level_dimensions[level][0]) + + @staticmethod + def get_downsample_ratio(wsi, level: int) -> float: + """ + Returns the down-sampling ratio of the whole slide image at a given level. + + Args: + wsi: a whole slide image object loaded from a file + level: the level number where the size is calculated + + """ + return wsi.level_downsamples[level] # type: ignore + + def get_file_path(self, wsi) -> str: + """Return the file path for the WSI object""" + return str(abspath(wsi._filename)) + + def read(self, data: Union[Sequence[PathLike], PathLike, np.ndarray], **kwargs): + """ + Read whole slide image objects from given file or list of files. + + Args: + data: file name or a list of file names to read. + kwargs: additional args that overrides `self.kwargs` for existing keys. + + Returns: + whole slide image object or list of such objects + + """ + wsi_list: List = [] + + filenames: Sequence[PathLike] = ensure_tuple(data) + kwargs_ = self.kwargs.copy() + kwargs_.update(kwargs) + for filename in filenames: + wsi = OpenSlide(filename, **kwargs_) + wsi_list.append(wsi) + + return wsi_list if len(filenames) > 1 else wsi_list[0] + + def get_patch( + self, wsi, location: Tuple[int, int], size: Tuple[int, int], level: int, dtype: DtypeLike, mode: str + ) -> np.ndarray: + """ + Extracts and returns a patch image form the whole slide image. + + Args: + wsi: a whole slide image object loaded from a file or a lis of such objects + location: (top, left) tuple giving the top left pixel in the level 0 reference frame. Defaults to (0, 0). + size: (height, width) tuple giving the patch size at the given level (`level`). + If None, it is set to the full image size at the given level. + level: the level number. Defaults to 0 + dtype: the data type of output image + mode: the output image mode, 'RGB' or 'RGBA' + + """ + # Extract a patch or the entire image + # (reverse the order of location and size to become WxH for OpenSlide) + pil_patch = wsi.read_region(location=location[::-1], size=size[::-1], level=level) + + # convert to RGB/RGBA + pil_patch = pil_patch.convert(mode) + + # Convert to numpy + patch = np.asarray(pil_patch, dtype=dtype) + + # Make the channel to desired dimensions + patch = np.moveaxis(patch, -1, self.channel_dim) + + return patch diff --git a/monai/engines/evaluator.py b/monai/engines/evaluator.py index c3e8c456b7..7999bb9bd6 100644 --- a/monai/engines/evaluator.py +++ b/monai/engines/evaluator.py @@ -9,7 +9,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence import torch from torch.utils.data import DataLoader @@ -62,7 +64,7 @@ class Evaluator(Workflow): it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. val_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision evaluation, default is False. mode: model forward mode during evaluation, should be 'eval' or 'train', which maps to `model.eval()` or `model.train()`, default to 'eval'. @@ -74,27 +76,33 @@ class Evaluator(Workflow): decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. """ def __init__( self, device: torch.device, - val_data_loader: Union[Iterable, DataLoader], - epoch_length: Optional[int] = None, + val_data_loader: Iterable | DataLoader, + epoch_length: int | None = None, non_blocking: bool = False, prepare_batch: Callable = default_prepare_batch, - iteration_update: Optional[Callable[[Engine, Any], Any]] = None, - postprocessing: Optional[Transform] = None, - key_val_metric: Optional[Dict[str, Metric]] = None, - additional_metrics: Optional[Dict[str, Metric]] = None, + iteration_update: Callable[[Engine, Any], Any] | None = None, + postprocessing: Transform | None = None, + key_val_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, metric_cmp_fn: Callable = default_metric_cmp_fn, - val_handlers: Optional[Sequence] = None, + val_handlers: Sequence | None = None, amp: bool = False, - mode: Union[ForwardMode, str] = ForwardMode.EVAL, - event_names: Optional[List[Union[str, EventEnum]]] = None, - event_to_attr: Optional[dict] = None, + mode: ForwardMode | str = ForwardMode.EVAL, + event_names: list[str | EventEnum] | None = None, + event_to_attr: dict | None = None, decollate: bool = True, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, ) -> None: super().__init__( device=device, @@ -113,6 +121,8 @@ def __init__( event_names=event_names, event_to_attr=event_to_attr, decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, ) mode = look_up_option(mode, ForwardMode) if mode == ForwardMode.EVAL: @@ -136,7 +146,7 @@ def run(self, global_epoch: int = 1) -> None: self.state.iteration = 0 super().run() - def get_validation_stats(self) -> Dict[str, float]: + def get_validation_stats(self) -> dict[str, float]: return {"best_validation_metric": self.state.best_metric, "best_validation_epoch": self.state.best_metric_epoch} @@ -169,7 +179,7 @@ class SupervisedEvaluator(Evaluator): it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. val_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision evaluation, default is False. mode: model forward mode during evaluation, should be 'eval' or 'train', which maps to `model.eval()` or `model.train()`, default to 'eval'. @@ -181,29 +191,35 @@ class SupervisedEvaluator(Evaluator): decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. """ def __init__( self, device: torch.device, - val_data_loader: Union[Iterable, DataLoader], + val_data_loader: Iterable | DataLoader, network: torch.nn.Module, - epoch_length: Optional[int] = None, + epoch_length: int | None = None, non_blocking: bool = False, prepare_batch: Callable = default_prepare_batch, - iteration_update: Optional[Callable[[Engine, Any], Any]] = None, - inferer: Optional[Inferer] = None, - postprocessing: Optional[Transform] = None, - key_val_metric: Optional[Dict[str, Metric]] = None, - additional_metrics: Optional[Dict[str, Metric]] = None, + iteration_update: Callable[[Engine, Any], Any] | None = None, + inferer: Inferer | None = None, + postprocessing: Transform | None = None, + key_val_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, metric_cmp_fn: Callable = default_metric_cmp_fn, - val_handlers: Optional[Sequence] = None, + val_handlers: Sequence | None = None, amp: bool = False, - mode: Union[ForwardMode, str] = ForwardMode.EVAL, - event_names: Optional[List[Union[str, EventEnum]]] = None, - event_to_attr: Optional[dict] = None, + mode: ForwardMode | str = ForwardMode.EVAL, + event_names: list[str | EventEnum] | None = None, + event_to_attr: dict | None = None, decollate: bool = True, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, ) -> None: super().__init__( device=device, @@ -222,12 +238,14 @@ def __init__( event_names=event_names, event_to_attr=event_to_attr, decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, ) self.network = network self.inferer = SimpleInferer() if inferer is None else inferer - def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): + def _iteration(self, engine: SupervisedEvaluator, batchdata: dict[str, torch.Tensor]): """ callback function for the Supervised Evaluation processing logic of 1 iteration in Ignite Engine. Return below items in a dictionary: @@ -236,7 +254,7 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): - PRED: prediction result of model. Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. + engine: `SupervisedEvaluator` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: @@ -245,24 +263,25 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): """ if batchdata is None: raise ValueError("Must provide batch data for current iteration.") - batch = self.prepare_batch(batchdata, engine.state.device, engine.non_blocking) # type: ignore + batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) if len(batch) == 2: inputs, targets = batch - args: Tuple = () - kwargs: Dict = {} + args: tuple = () + kwargs: dict = {} else: inputs, targets, args, kwargs = batch # put iteration outputs into engine.state - engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} # type: ignore + engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} # execute forward computation - with self.mode(self.network): - if self.amp: - with torch.cuda.amp.autocast(): - engine.state.output[Keys.PRED] = self.inferer(inputs, self.network, *args, **kwargs) # type: ignore + with engine.mode(engine.network): + + if engine.amp: + with torch.cuda.amp.autocast(**engine.amp_kwargs): + engine.state.output[Keys.PRED] = engine.inferer(inputs, engine.network, *args, **kwargs) else: - engine.state.output[Keys.PRED] = self.inferer(inputs, self.network, *args, **kwargs) # type: ignore + engine.state.output[Keys.PRED] = engine.inferer(inputs, engine.network, *args, **kwargs) engine.fire_event(IterationEvents.FORWARD_COMPLETED) engine.fire_event(IterationEvents.MODEL_COMPLETED) @@ -302,7 +321,7 @@ class EnsembleEvaluator(Evaluator): it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. val_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision evaluation, default is False. mode: model forward mode during evaluation, should be 'eval' or 'train', which maps to `model.eval()` or `model.train()`, default to 'eval'. @@ -314,30 +333,36 @@ class EnsembleEvaluator(Evaluator): decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. """ def __init__( self, device: torch.device, - val_data_loader: Union[Iterable, DataLoader], + val_data_loader: Iterable | DataLoader, networks: Sequence[torch.nn.Module], - pred_keys: Optional[KeysCollection] = None, - epoch_length: Optional[int] = None, + pred_keys: KeysCollection | None = None, + epoch_length: int | None = None, non_blocking: bool = False, prepare_batch: Callable = default_prepare_batch, - iteration_update: Optional[Callable[[Engine, Any], Any]] = None, - inferer: Optional[Inferer] = None, - postprocessing: Optional[Transform] = None, - key_val_metric: Optional[Dict[str, Metric]] = None, - additional_metrics: Optional[Dict[str, Metric]] = None, + iteration_update: Callable[[Engine, Any], Any] | None = None, + inferer: Inferer | None = None, + postprocessing: Transform | None = None, + key_val_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, metric_cmp_fn: Callable = default_metric_cmp_fn, - val_handlers: Optional[Sequence] = None, + val_handlers: Sequence | None = None, amp: bool = False, - mode: Union[ForwardMode, str] = ForwardMode.EVAL, - event_names: Optional[List[Union[str, EventEnum]]] = None, - event_to_attr: Optional[dict] = None, + mode: ForwardMode | str = ForwardMode.EVAL, + event_names: list[str | EventEnum] | None = None, + event_to_attr: dict | None = None, decollate: bool = True, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, ) -> None: super().__init__( device=device, @@ -356,6 +381,8 @@ def __init__( event_names=event_names, event_to_attr=event_to_attr, decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, ) self.networks = ensure_tuple(networks) @@ -366,7 +393,7 @@ def __init__( raise ValueError("length of `pred_keys` must be same as the length of `networks`.") self.inferer = SimpleInferer() if inferer is None else inferer - def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): + def _iteration(self, engine: EnsembleEvaluator, batchdata: dict[str, torch.Tensor]): """ callback function for the Supervised Evaluation processing logic of 1 iteration in Ignite Engine. Return below items in a dictionary: @@ -378,7 +405,7 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): - pred_keys[N]: prediction result of network N. Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. + engine: `EnsembleEvaluator` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: @@ -387,29 +414,29 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): """ if batchdata is None: raise ValueError("Must provide batch data for current iteration.") - batch = self.prepare_batch(batchdata, engine.state.device, engine.non_blocking) # type: ignore + batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) if len(batch) == 2: inputs, targets = batch - args: Tuple = () - kwargs: Dict = {} + args: tuple = () + kwargs: dict = {} else: inputs, targets, args, kwargs = batch # put iteration outputs into engine.state - engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} # type: ignore + engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} - for idx, network in enumerate(self.networks): - with self.mode(network): - if self.amp: - with torch.cuda.amp.autocast(): + for idx, network in enumerate(engine.networks): + with engine.mode(network): + if engine.amp: + with torch.cuda.amp.autocast(**engine.amp_kwargs): if isinstance(engine.state.output, dict): engine.state.output.update( - {self.pred_keys[idx]: self.inferer(inputs, network, *args, **kwargs)} + {engine.pred_keys[idx]: engine.inferer(inputs, network, *args, **kwargs)} ) else: if isinstance(engine.state.output, dict): engine.state.output.update( - {self.pred_keys[idx]: self.inferer(inputs, network, *args, **kwargs)} + {engine.pred_keys[idx]: engine.inferer(inputs, network, *args, **kwargs)} ) engine.fire_event(IterationEvents.FORWARD_COMPLETED) engine.fire_event(IterationEvents.MODEL_COMPLETED) diff --git a/monai/engines/trainer.py b/monai/engines/trainer.py index 774e535e7f..2b7a1acd2a 100644 --- a/monai/engines/trainer.py +++ b/monai/engines/trainer.py @@ -9,7 +9,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence import torch from torch.optim.optimizer import Optimizer @@ -26,7 +28,7 @@ from monai.engines.workflow import Workflow from monai.inferers import Inferer, SimpleInferer from monai.transforms import Transform -from monai.utils import min_version, optional_import, pytorch_after +from monai.utils import min_version, optional_import from monai.utils.enums import CommonKeys as Keys if TYPE_CHECKING: @@ -55,7 +57,7 @@ def run(self) -> None: self.scaler = torch.cuda.amp.GradScaler() if self.amp else None super().run() - def get_train_stats(self) -> Dict[str, float]: + def get_train_stats(self) -> dict[str, float]: return {"total_epochs": self.state.max_epochs, "total_iterations": self.state.epoch_length} @@ -93,7 +95,7 @@ class SupervisedTrainer(Trainer): it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision training, default is False. event_names: additional custom ignite events that will register to the engine. new events can be a list of str or `ignite.engine.events.EventEnum`. @@ -105,6 +107,10 @@ class SupervisedTrainer(Trainer): default to `True`. optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. """ @@ -112,25 +118,27 @@ def __init__( self, device: torch.device, max_epochs: int, - train_data_loader: Union[Iterable, DataLoader], + train_data_loader: Iterable | DataLoader, network: torch.nn.Module, optimizer: Optimizer, loss_function: Callable, - epoch_length: Optional[int] = None, + epoch_length: int | None = None, non_blocking: bool = False, prepare_batch: Callable = default_prepare_batch, - iteration_update: Optional[Callable[[Engine, Any], Any]] = None, - inferer: Optional[Inferer] = None, - postprocessing: Optional[Transform] = None, - key_train_metric: Optional[Dict[str, Metric]] = None, - additional_metrics: Optional[Dict[str, Metric]] = None, + iteration_update: Callable[[Engine, Any], Any] | None = None, + inferer: Inferer | None = None, + postprocessing: Transform | None = None, + key_train_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, metric_cmp_fn: Callable = default_metric_cmp_fn, - train_handlers: Optional[Sequence] = None, + train_handlers: Sequence | None = None, amp: bool = False, - event_names: Optional[List[Union[str, EventEnum]]] = None, - event_to_attr: Optional[dict] = None, + event_names: list[str | EventEnum] | None = None, + event_to_attr: dict | None = None, decollate: bool = True, optim_set_to_none: bool = False, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, ) -> None: super().__init__( device=device, @@ -149,6 +157,8 @@ def __init__( event_names=event_names, event_to_attr=event_to_attr, decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, ) self.network = network @@ -157,7 +167,7 @@ def __init__( self.inferer = SimpleInferer() if inferer is None else inferer self.optim_set_to_none = optim_set_to_none - def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): + def _iteration(self, engine: SupervisedTrainer, batchdata: dict[str, torch.Tensor]): """ Callback function for the Supervised Training processing logic of 1 iteration in Ignite Engine. Return below items in a dictionary: @@ -167,7 +177,7 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): - LOSS: loss value computed by loss function. Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. + engine: `SupervisedTrainer` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: @@ -176,41 +186,37 @@ def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): """ if batchdata is None: raise ValueError("Must provide batch data for current iteration.") - batch = self.prepare_batch(batchdata, engine.state.device, engine.non_blocking) # type: ignore + batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) if len(batch) == 2: inputs, targets = batch - args: Tuple = () - kwargs: Dict = {} + args: tuple = () + kwargs: dict = {} else: inputs, targets, args, kwargs = batch # put iteration outputs into engine.state - engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} # type: ignore + engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets} def _compute_pred_loss(): - engine.state.output[Keys.PRED] = self.inferer(inputs, self.network, *args, **kwargs) + engine.state.output[Keys.PRED] = engine.inferer(inputs, engine.network, *args, **kwargs) engine.fire_event(IterationEvents.FORWARD_COMPLETED) - engine.state.output[Keys.LOSS] = self.loss_function(engine.state.output[Keys.PRED], targets).mean() + engine.state.output[Keys.LOSS] = engine.loss_function(engine.state.output[Keys.PRED], targets).mean() engine.fire_event(IterationEvents.LOSS_COMPLETED) - self.network.train() - # `set_to_none` only work from PyTorch 1.7.0 - if not pytorch_after(1, 7): - self.optimizer.zero_grad() - else: - self.optimizer.zero_grad(set_to_none=self.optim_set_to_none) + engine.network.train() + engine.optimizer.zero_grad(set_to_none=engine.optim_set_to_none) - if self.amp and self.scaler is not None: - with torch.cuda.amp.autocast(): + if engine.amp and engine.scaler is not None: + with torch.cuda.amp.autocast(**engine.amp_kwargs): _compute_pred_loss() - self.scaler.scale(engine.state.output[Keys.LOSS]).backward() # type: ignore + engine.scaler.scale(engine.state.output[Keys.LOSS]).backward() engine.fire_event(IterationEvents.BACKWARD_COMPLETED) - self.scaler.step(self.optimizer) - self.scaler.update() + engine.scaler.step(engine.optimizer) + engine.scaler.update() else: _compute_pred_loss() - engine.state.output[Keys.LOSS].backward() # type: ignore + engine.state.output[Keys.LOSS].backward() engine.fire_event(IterationEvents.BACKWARD_COMPLETED) - self.optimizer.step() + engine.optimizer.step() engine.fire_event(IterationEvents.MODEL_COMPLETED) return engine.state.output @@ -265,12 +271,16 @@ class GanTrainer(Trainer): it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. """ @@ -285,23 +295,25 @@ def __init__( d_network: torch.nn.Module, d_optimizer: Optimizer, d_loss_function: Callable, - epoch_length: Optional[int] = None, - g_inferer: Optional[Inferer] = None, - d_inferer: Optional[Inferer] = None, + epoch_length: int | None = None, + g_inferer: Inferer | None = None, + d_inferer: Inferer | None = None, d_train_steps: int = 1, latent_shape: int = 64, non_blocking: bool = False, d_prepare_batch: Callable = default_prepare_batch, g_prepare_batch: Callable = default_make_latent, g_update_latents: bool = True, - iteration_update: Optional[Callable[[Engine, Any], Any]] = None, - postprocessing: Optional[Transform] = None, - key_train_metric: Optional[Dict[str, Metric]] = None, - additional_metrics: Optional[Dict[str, Metric]] = None, + iteration_update: Callable[[Engine, Any], Any] | None = None, + postprocessing: Transform | None = None, + key_train_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, metric_cmp_fn: Callable = default_metric_cmp_fn, - train_handlers: Optional[Sequence] = None, + train_handlers: Sequence | None = None, decollate: bool = True, optim_set_to_none: bool = False, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, ): if not isinstance(train_data_loader, DataLoader): raise ValueError("train_data_loader must be PyTorch DataLoader.") @@ -321,6 +333,8 @@ def __init__( handlers=train_handlers, postprocessing=postprocessing, decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, ) self.g_network = g_network self.g_optimizer = g_optimizer @@ -337,13 +351,13 @@ def __init__( self.optim_set_to_none = optim_set_to_none def _iteration( - self, engine: Engine, batchdata: Union[Dict, Sequence] - ) -> Dict[str, Union[torch.Tensor, int, float, bool]]: + self, engine: GanTrainer, batchdata: dict | Sequence + ) -> dict[str, torch.Tensor | int | float | bool]: """ Callback function for Adversarial Training processing logic of 1 iteration in Ignite Engine. Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. + engine: `GanTrainer` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: @@ -353,45 +367,40 @@ def _iteration( if batchdata is None: raise ValueError("must provide batch data for current iteration.") - d_input = self.prepare_batch(batchdata, engine.state.device, engine.non_blocking) # type: ignore - batch_size = self.data_loader.batch_size # type: ignore - g_input = self.g_prepare_batch( + d_input = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) + batch_size = engine.data_loader.batch_size # type: ignore + g_input = engine.g_prepare_batch( num_latents=batch_size, - latent_size=self.latent_shape, - device=engine.state.device, # type: ignore - non_blocking=engine.non_blocking, # type: ignore + latent_size=engine.latent_shape, + device=engine.state.device, + non_blocking=engine.non_blocking, + **engine.to_kwargs, ) - g_output = self.g_inferer(g_input, self.g_network) + g_output = engine.g_inferer(g_input, engine.g_network) # Train Discriminator d_total_loss = torch.zeros(1) - for _ in range(self.d_train_steps): - # `set_to_none` only work from PyTorch 1.7.0 - if not pytorch_after(1, 7): - self.d_optimizer.zero_grad() - else: - self.d_optimizer.zero_grad(set_to_none=self.optim_set_to_none) - dloss = self.d_loss_function(g_output, d_input) + for _ in range(engine.d_train_steps): + engine.d_optimizer.zero_grad(set_to_none=engine.optim_set_to_none) + dloss = engine.d_loss_function(g_output, d_input) dloss.backward() - self.d_optimizer.step() + engine.d_optimizer.step() d_total_loss += dloss.item() # Train Generator - if self.g_update_latents: - g_input = self.g_prepare_batch( + if engine.g_update_latents: + g_input = engine.g_prepare_batch( num_latents=batch_size, - latent_size=self.latent_shape, - device=engine.state.device, # type: ignore - non_blocking=engine.non_blocking, # type: ignore + latent_size=engine.latent_shape, + device=engine.state.device, + non_blocking=engine.non_blocking, + **engine.to_kwargs, ) - g_output = self.g_inferer(g_input, self.g_network) - if not pytorch_after(1, 7): - self.g_optimizer.zero_grad() - else: - self.g_optimizer.zero_grad(set_to_none=self.optim_set_to_none) - g_loss = self.g_loss_function(g_output) + g_output = engine.g_inferer(g_input, engine.g_network) + engine.g_optimizer.zero_grad(set_to_none=engine.optim_set_to_none) + g_loss = engine.g_loss_function(g_output) g_loss.backward() - self.g_optimizer.step() + engine.g_optimizer.step() return { GanKeys.REALS: d_input, diff --git a/monai/engines/utils.py b/monai/engines/utils.py index 726dfc8e98..8f3a57beda 100644 --- a/monai/engines/utils.py +++ b/monai/engines/utils.py @@ -104,12 +104,16 @@ def get_devices_spec(devices: Optional[Sequence[torch.device]] = None) -> List[t def default_prepare_batch( - batchdata: Dict[str, torch.Tensor], device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False + batchdata: Dict[str, torch.Tensor], + device: Optional[Union[str, torch.device]] = None, + non_blocking: bool = False, + **kwargs, ) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]: """ Default function to prepare the data for current iteration. - Refer to ignite: https://pytorch.org/ignite/v0.4.5/generated/ignite.engine.create_supervised_trainer.html - #ignite.engine.create_supervised_trainer. + Args `batchdata`, `device`, `non_blocking` refer to the ignite API: + https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. + `kwargs` supports other args for `Tensor.to()` API. Returns: image, label(optional). @@ -119,18 +123,21 @@ def default_prepare_batch( raise AssertionError("default prepare_batch expects dictionary input data.") if isinstance(batchdata.get(CommonKeys.LABEL), torch.Tensor): return ( - batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking), - batchdata[CommonKeys.LABEL].to(device=device, non_blocking=non_blocking), + batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking, **kwargs), + batchdata[CommonKeys.LABEL].to(device=device, non_blocking=non_blocking, **kwargs), ) if GanKeys.REALS in batchdata: - return batchdata[GanKeys.REALS].to(device=device, non_blocking=non_blocking) - return batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking), None + return batchdata[GanKeys.REALS].to(device=device, non_blocking=non_blocking, **kwargs) + return batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking, **kwargs), None class PrepareBatch(ABC): """ Interface of customized prepare_batch in the trainer or evaluator workflows. It takes the data of current batch, target device and non_blocking flag as input. + Args `batchdata`, `device`, `non_blocking` refer to the ignite API: + https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. + `kwargs` supports other args for `Tensor.to()` API. """ @@ -140,6 +147,7 @@ def __call__( batchdata: Dict[str, torch.Tensor], device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False, + **kwargs, ): raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") @@ -155,8 +163,15 @@ def __call__( batchdata: Dict[str, torch.Tensor], device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False, + **kwargs, ): - return default_prepare_batch(batchdata, device, non_blocking) + """ + Args `batchdata`, `device`, `non_blocking` refer to the ignite API: + https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. + `kwargs` supports other args for `Tensor.to()` API. + + """ + return default_prepare_batch(batchdata, device, non_blocking, **kwargs) class PrepareBatchExtraInput(PrepareBatch): @@ -181,29 +196,42 @@ def __call__( batchdata: Dict[str, torch.Tensor], device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False, + **kwargs, ): - image, label = default_prepare_batch(batchdata, device, non_blocking) - args = list() - kwargs = dict() + """ + Args `batchdata`, `device`, `non_blocking` refer to the ignite API: + https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. + `kwargs` supports other args for `Tensor.to()` API. + + """ + image, label = default_prepare_batch(batchdata, device, non_blocking, **kwargs) + args_ = list() + kwargs_ = dict() def _get_data(key: str): data = batchdata[key] - return data.to(device=device, non_blocking=non_blocking) if isinstance(data, torch.Tensor) else data + return ( + data.to(device=device, non_blocking=non_blocking, **kwargs) if isinstance(data, torch.Tensor) else data + ) if isinstance(self.extra_keys, (str, list, tuple)): for k in ensure_tuple(self.extra_keys): - args.append(_get_data(k)) + args_.append(_get_data(k)) elif isinstance(self.extra_keys, dict): for k, v in self.extra_keys.items(): - kwargs.update({k: _get_data(v)}) + kwargs_.update({k: _get_data(v)}) - return image, label, tuple(args), kwargs + return image, label, tuple(args_), kwargs_ def default_make_latent( - num_latents: int, latent_size: int, device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False + num_latents: int, + latent_size: int, + device: Optional[Union[str, torch.device]] = None, + non_blocking: bool = False, + **kwargs, ) -> torch.Tensor: - return torch.randn(num_latents, latent_size).to(device=device, non_blocking=non_blocking) + return torch.randn(num_latents, latent_size).to(device=device, non_blocking=non_blocking, **kwargs) def engine_apply_transform(batch: Any, output: Any, transform: Callable[..., Dict]): diff --git a/monai/engines/workflow.py b/monai/engines/workflow.py index 65bb313e53..8349ff82ab 100644 --- a/monai/engines/workflow.py +++ b/monai/engines/workflow.py @@ -84,7 +84,7 @@ class Workflow(IgniteEngine): # type: ignore[valid-type, misc] # due to optiona it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: - CheckpointHandler, StatsHandler, SegmentationSaver, etc. + CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision training or inference, default is False. event_names: additional custom ignite events that will register to the engine. new events can be a list of str or `ignite.engine.events.EventEnum`. @@ -94,6 +94,10 @@ class Workflow(IgniteEngine): # type: ignore[valid-type, misc] # due to optiona decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. Raises: TypeError: When ``device`` is not a ``torch.Device``. @@ -121,6 +125,8 @@ def __init__( event_names: Optional[List[Union[str, EventEnum]]] = None, event_to_attr: Optional[dict] = None, decollate: bool = True, + to_kwargs: Optional[Dict] = None, + amp_kwargs: Optional[Dict] = None, ) -> None: if iteration_update is not None: super().__init__(iteration_update) @@ -166,6 +172,8 @@ def set_sampler_epoch(engine: Engine): self.prepare_batch = prepare_batch self.metric_cmp_fn = metric_cmp_fn self.amp = amp + self.to_kwargs = {} if to_kwargs is None else to_kwargs + self.amp_kwargs = {} if amp_kwargs is None else amp_kwargs self.scaler: Optional[torch.cuda.amp.GradScaler] = None if event_names is None: @@ -242,8 +250,8 @@ def _register_metrics(self, k_metric: Dict, add_metrics: Optional[Dict] = None): metric.attach(self, name) @self.on(Events.EPOCH_COMPLETED) - def _compare_metrics(engine: Engine) -> None: - key_metric_name = engine.state.key_metric_name # type: ignore + def _compare_metrics(engine: Workflow) -> None: + key_metric_name = engine.state.key_metric_name if key_metric_name is not None: current_val_metric = engine.state.metrics[key_metric_name] if not is_scalar(current_val_metric): @@ -253,10 +261,10 @@ def _compare_metrics(engine: Engine) -> None: ) return - if self.metric_cmp_fn(current_val_metric, engine.state.best_metric): # type: ignore + if self.metric_cmp_fn(current_val_metric, engine.state.best_metric): self.logger.info(f"Got new best metric of {key_metric_name}: {current_val_metric}") - engine.state.best_metric = current_val_metric # type: ignore - engine.state.best_metric_epoch = engine.state.epoch # type: ignore + engine.state.best_metric = current_val_metric + engine.state.best_metric_epoch = engine.state.epoch def _register_handlers(self, handlers: Sequence): """ @@ -281,7 +289,7 @@ def run(self) -> None: return super().run(data=self.data_loader, max_epochs=self.state.max_epochs) - def _iteration(self, engine: Engine, batchdata: Dict[str, torch.Tensor]): + def _iteration(self, engine, batchdata: Dict[str, torch.Tensor]): """ Abstract callback function for the processing logic of 1 iteration in Ignite Engine. Need subclass to implement different logics, like SupervisedTrainer/Evaluator, GANTrainer, etc. diff --git a/monai/handlers/__init__.py b/monai/handlers/__init__.py index 649cc3cae6..cffbe46391 100644 --- a/monai/handlers/__init__.py +++ b/monai/handlers/__init__.py @@ -26,9 +26,9 @@ from .nvtx_handlers import MarkHandler, RangeHandler, RangePopHandler, RangePushHandler from .parameter_scheduler import ParamSchedulerHandler from .postprocessing import PostProcessing +from .probability_maps import ProbMapProducer from .regression_metrics import MeanAbsoluteError, MeanSquaredError, PeakSignalToNoiseRatio, RootMeanSquaredError from .roc_auc import ROCAUC -from .segmentation_saver import SegmentationSaver from .smartcache_handler import SmartCacheHandler from .stats_handler import StatsHandler from .surface_distance import SurfaceDistance diff --git a/monai/handlers/probability_maps.py b/monai/handlers/probability_maps.py new file mode 100644 index 0000000000..5c15704d50 --- /dev/null +++ b/monai/handlers/probability_maps.py @@ -0,0 +1,121 @@ +# 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. + +import logging +import os +import threading +from typing import TYPE_CHECKING, Dict, Optional + +import numpy as np + +from monai.config import DtypeLike, IgniteInfo +from monai.utils import ProbMapKeys, min_version, optional_import + +Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") +if TYPE_CHECKING: + from ignite.engine import Engine +else: + Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") + + +class ProbMapProducer: + """ + Event handler triggered on completing every iteration to calculate and save the probability map + """ + + def __init__( + self, + output_dir: str = "./", + output_postfix: str = "", + prob_key: str = "pred", + dtype: DtypeLike = np.float64, + name: Optional[str] = None, + ) -> None: + """ + Args: + output_dir: output directory to save probability maps. + output_postfix: a string appended to all output file names. + prob_key: the key associated to the probability output of the model + dtype: the data type in which the probability map is stored. Default np.float64. + name: identifier of logging.logger to use, defaulting to `engine.logger`. + + """ + self.logger = logging.getLogger(name) + self._name = name + self.output_dir = output_dir + self.output_postfix = output_postfix + self.prob_key = prob_key + self.dtype = dtype + self.prob_map: Dict[str, np.ndarray] = {} + self.counter: Dict[str, int] = {} + self.num_done_images: int = 0 + self.num_images: int = 0 + self.lock = threading.Lock() + + def attach(self, engine: Engine) -> None: + """ + Args: + engine: Ignite Engine, it can be a trainer, validator or evaluator. + """ + + image_data = engine.data_loader.dataset.image_data # type: ignore + self.num_images = len(image_data) + + # Initialized probability maps for all the images + for sample in image_data: + name = sample[ProbMapKeys.NAME] + self.counter[name] = sample[ProbMapKeys.COUNT] + self.prob_map[name] = np.zeros(sample[ProbMapKeys.SIZE], dtype=self.dtype) + + if self._name is None: + self.logger = engine.logger + if not engine.has_event_handler(self, Events.ITERATION_COMPLETED): + engine.add_event_handler(Events.ITERATION_COMPLETED, self) + if not engine.has_event_handler(self.finalize, Events.COMPLETED): + engine.add_event_handler(Events.COMPLETED, self.finalize) + + def __call__(self, engine: Engine) -> None: + """ + This method assumes self.batch_transform will extract metadata from the input batch. + + Args: + engine: Ignite Engine, it can be a trainer, validator or evaluator. + """ + if not isinstance(engine.state.batch, dict) or not isinstance(engine.state.output, dict): + raise ValueError("engine.state.batch and engine.state.output must be dictionaries.") + names = engine.state.batch["metadata"][ProbMapKeys.NAME] + locs = engine.state.batch["metadata"][ProbMapKeys.LOCATION] + probs = engine.state.output[self.prob_key] + for name, loc, prob in zip(names, locs, probs): + self.prob_map[name][loc] = prob + with self.lock: + self.counter[name] -= 1 + if self.counter[name] == 0: + self.save_prob_map(name) + + def save_prob_map(self, name: str) -> None: + """ + This method save the probability map for an image, when its inference is finished, + and delete that probability map from memory. + + Args: + name: the name of image to be saved. + """ + file_path = os.path.join(self.output_dir, name) + np.save(file_path + self.output_postfix + ".npy", self.prob_map[name]) + + self.num_done_images += 1 + self.logger.info(f"Inference of '{name}' is done [{self.num_done_images}/{self.num_images}]!") + del self.prob_map[name] + del self.counter[name] + + def finalize(self, engine: Engine): + self.logger.info(f"Probability map is created for {self.num_done_images}/{self.num_images} images!") diff --git a/monai/handlers/segmentation_saver.py b/monai/handlers/segmentation_saver.py deleted file mode 100644 index 40bb5f8bed..0000000000 --- a/monai/handlers/segmentation_saver.py +++ /dev/null @@ -1,173 +0,0 @@ -# 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. - -import logging -from typing import TYPE_CHECKING, Callable, Optional, Union - -import numpy as np - -from monai.config import DtypeLike, IgniteInfo -from monai.data import decollate_batch -from monai.transforms import SaveImage -from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, deprecated, min_version, optional_import - -Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") -if TYPE_CHECKING: - from ignite.engine import Engine -else: - Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") - - -@deprecated(since="0.6.0", removed="0.9.0", msg_suffix="Please consider using `SaveImage[d]` transform instead.") -class SegmentationSaver: - """ - Event handler triggered on completing every iteration to save the segmentation predictions into files. - It can extract the input image meta data(filename, affine, original_shape, etc.) and resample the predictions - based on the meta data. - The name of saved file will be `{input_image_name}_{output_postfix}{output_ext}`, - where the input image name is extracted from the meta data dictionary. If no meta data provided, - use index from 0 as the filename prefix. - The predictions can be PyTorch Tensor with [B, C, H, W, [D]] shape or a list of Tensor without batch dim. - - .. deprecated:: 0.6.0 - Use :class:`monai.transforms.SaveImage` or :class:`monai.transforms.SaveImaged` instead. - - """ - - def __init__( - self, - output_dir: str = "./", - output_postfix: str = "seg", - output_ext: str = ".nii.gz", - resample: bool = True, - mode: Union[GridSampleMode, InterpolateMode, str] = "nearest", - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, - scale: Optional[int] = None, - dtype: DtypeLike = np.float64, - output_dtype: DtypeLike = np.float32, - squeeze_end_dims: bool = True, - data_root_dir: str = "", - separate_folder: bool = True, - batch_transform: Callable = lambda x: x, - output_transform: Callable = lambda x: x, - name: Optional[str] = None, - ) -> None: - """ - Args: - output_dir: output image directory. - output_postfix: a string appended to all output file names, default to `seg`. - output_ext: output file extension name, available extensions: `.nii.gz`, `.nii`, `.png`. - resample: whether to resample before saving the data array. - if saving PNG format image, based on the `spatial_shape` from metadata. - if saving NIfTI format image, based on the `original_affine` from metadata. - mode: This option is used when ``resample = True``. Defaults to ``"nearest"``. - - - NIfTI files {``"bilinear"``, ``"nearest"``} - Interpolation mode to calculate output values. - See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html - - PNG files {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} - The interpolation mode. - See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html - - padding_mode: This option is used when ``resample = True``. Defaults to ``"border"``. - - - NIfTI files {``"zeros"``, ``"border"``, ``"reflection"``} - Padding mode for outside grid values. - See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html - - PNG files - This option is ignored. - - scale: {``255``, ``65535``} postprocess data by clipping to [0, 1] and scaling - [0, 255] (uint8) or [0, 65535] (uint16). Default is None to disable scaling. - It's used for PNG format only. - dtype: data type for resampling computation. Defaults to ``np.float64`` for best precision. - If None, use the data type of input data. - It's used for Nifti format only. - output_dtype: data type for saving data. Defaults to ``np.float32``, it's used for Nifti format only. - squeeze_end_dims: if True, any trailing singleton dimensions will be removed (after the channel - has been moved to the end). So if input is (C,H,W,D), this will be altered to (H,W,D,C), and - then if C==1, it will be saved as (H,W,D). If D also ==1, it will be saved as (H,W). If false, - image will always be saved as (H,W,D,C). - it's used for NIfTI format only. - data_root_dir: if not empty, it specifies the beginning parts of the input file's - absolute path. it's used to compute `input_file_rel_path`, the relative path to the file from - `data_root_dir` to preserve folder structure when saving in case there are files in different - folders with the same file names. for example: - input_file_name: /foo/bar/test1/image.nii, - output_postfix: seg - output_ext: nii.gz - output_dir: /output, - data_root_dir: /foo/bar, - output will be: /output/test1/image/image_seg.nii.gz - separate_folder: whether to save every file in a separate folder, for example: if input filename is - `image.nii`, postfix is `seg` and folder_path is `output`, if `True`, save as: - `output/image/image_seg.nii`, if `False`, save as `output/image_seg.nii`. default to `True`. - batch_transform: a callable that is used to extract the `meta_data` dictionary of the input images - from `ignite.engine.state.batch`. the purpose is to extract necessary information from the meta data: - filename, affine, original_shape, etc. - `engine.state` and `batch_transform` inherit from the ignite concept: - https://pytorch.org/ignite/concepts.html#state, explanation and usage example are in the tutorial: - https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb. - output_transform: a callable that is used to extract the model prediction data from - `ignite.engine.state.output`. the first dimension of its output will be treated as the batch dimension. - each item in the batch will be saved individually. - `engine.state` and `output_transform` inherit from the ignite concept: - https://pytorch.org/ignite/concepts.html#state, explanation and usage example are in the tutorial: - https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb. - name: identifier of logging.logger to use, defaulting to `engine.logger`. - - """ - self._saver = SaveImage( - output_dir=output_dir, - output_postfix=output_postfix, - output_ext=output_ext, - resample=resample, - mode=mode, - padding_mode=padding_mode, - scale=scale, - dtype=dtype, - output_dtype=output_dtype, - squeeze_end_dims=squeeze_end_dims, - data_root_dir=data_root_dir, - separate_folder=separate_folder, - ) - self.batch_transform = batch_transform - self.output_transform = output_transform - - self.logger = logging.getLogger(name) - self._name = name - - def attach(self, engine: Engine) -> None: - """ - Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. - """ - if self._name is None: - self.logger = engine.logger - if not engine.has_event_handler(self, Events.ITERATION_COMPLETED): - engine.add_event_handler(Events.ITERATION_COMPLETED, self) - - def __call__(self, engine: Engine) -> None: - """ - This method assumes self.batch_transform will extract metadata from the input batch. - Output file datatype is determined from ``engine.state.output.dtype``. - - Args: - engine: Ignite Engine, it can be a trainer, validator or evaluator. - """ - meta_data = self.batch_transform(engine.state.batch) - if isinstance(meta_data, dict): - # decollate the `dictionary of list` to `list of dictionaries` - meta_data = decollate_batch(meta_data) - engine_output = self.output_transform(engine.state.output) - for m, o in zip(meta_data, engine_output): - self._saver(o, m) - self.logger.info("model outputs saved into files.") diff --git a/monai/handlers/tensorboard_handlers.py b/monai/handlers/tensorboard_handlers.py index 0105e7c8ca..445e3e76ca 100644 --- a/monai/handlers/tensorboard_handlers.py +++ b/monai/handlers/tensorboard_handlers.py @@ -20,13 +20,12 @@ from monai.visualize import plot_2d_or_3d_image Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") +SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") if TYPE_CHECKING: from ignite.engine import Engine - from torch.utils.tensorboard import SummaryWriter else: Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") - SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") DEFAULT_TAG = "Loss" @@ -42,7 +41,7 @@ class TensorBoardHandler: """ - def __init__(self, summary_writer: Optional[SummaryWriter] = None, log_dir: str = "./runs"): + def __init__(self, summary_writer=None, log_dir: str = "./runs"): if summary_writer is None: self._writer = SummaryWriter(log_dir=log_dir) self.internal_writer = True @@ -82,13 +81,13 @@ class TensorBoardStatsHandler(TensorBoardHandler): def __init__( self, - summary_writer: Optional[SummaryWriter] = None, + summary_writer=None, log_dir: str = "./runs", iteration_log: bool = True, epoch_log: bool = True, - epoch_event_writer: Optional[Callable[[Engine, SummaryWriter], Any]] = None, + epoch_event_writer: Optional[Callable[[Engine, Any], Any]] = None, epoch_interval: int = 1, - iteration_event_writer: Optional[Callable[[Engine, SummaryWriter], Any]] = None, + iteration_event_writer: Optional[Callable[[Engine, Any], Any]] = None, iteration_interval: int = 1, output_transform: Callable = lambda x: x[0], global_epoch_transform: Callable = lambda x: x, @@ -179,7 +178,7 @@ def iteration_completed(self, engine: Engine) -> None: else: self._default_iteration_writer(engine, self._writer) - def _write_scalar(self, _engine: Engine, writer: SummaryWriter, tag: str, value: Any, step: int) -> None: + def _write_scalar(self, _engine: Engine, writer, tag: str, value: Any, step: int) -> None: """ Write scale value into TensorBoard. Default to call `SummaryWriter.add_scalar()`. @@ -194,7 +193,7 @@ def _write_scalar(self, _engine: Engine, writer: SummaryWriter, tag: str, value: """ writer.add_scalar(tag, value, step) - def _default_epoch_writer(self, engine: Engine, writer: SummaryWriter) -> None: + def _default_epoch_writer(self, engine: Engine, writer) -> None: """ Execute epoch level event write operation. Default to write the values from Ignite `engine.state.metrics` dict and @@ -216,7 +215,7 @@ def _default_epoch_writer(self, engine: Engine, writer: SummaryWriter) -> None: self._write_scalar(engine, writer, attr, getattr(engine.state, attr, None), current_epoch) writer.flush() - def _default_iteration_writer(self, engine: Engine, writer: SummaryWriter) -> None: + def _default_iteration_writer(self, engine: Engine, writer) -> None: """ Execute iteration level event write operation based on Ignite `engine.state.output` data. Extract the values from `self.output_transform(engine.state.output)`. @@ -295,7 +294,7 @@ class TensorBoardImageHandler(TensorBoardHandler): def __init__( self, - summary_writer: Optional[SummaryWriter] = None, + summary_writer=None, log_dir: str = "./runs", interval: int = 1, epoch_level: bool = True, diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py index 331637ba94..084b1021c2 100644 --- a/monai/inferers/inferer.py +++ b/monai/inferers/inferer.py @@ -9,13 +9,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from abc import ABC, abstractmethod -from typing import Any, Callable, Optional, Sequence, Union +from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union import torch import torch.nn as nn -from monai.inferers.utils import sliding_window_inference +from monai.inferers.utils import compute_importance_map, sliding_window_inference from monai.utils import BlendMode, PytorchPadMode, ensure_tuple from monai.visualize import CAM, GradCAM, GradCAMpp @@ -120,6 +121,7 @@ class SlidingWindowInferer(Inferer): set to device=torch.device('cpu') the gpu memory consumption is less and independent of the `inputs` and `roi_size`. Output is on the `device`. progress: whether to print a tqdm progress bar. + cache_roi_weight_map: whether to precompute the ROI weight map. Note: ``sw_batch_size`` denotes the max number of windows per network inference iteration, @@ -139,6 +141,7 @@ def __init__( sw_device: Union[torch.device, str, None] = None, device: Union[torch.device, str, None] = None, progress: bool = False, + cache_roi_weight_map: bool = False, ) -> None: Inferer.__init__(self) self.roi_size = roi_size @@ -152,9 +155,30 @@ def __init__( self.device = device self.progress = progress + # compute_importance_map takes long time when computing on cpu. We thus + # compute it once if it's static and then save it for future usage + self.roi_weight_map = None + try: + if cache_roi_weight_map and isinstance(roi_size, Sequence) and min(roi_size) > 0: # non-dynamic roi size + if device is None: + device = "cpu" + self.roi_weight_map = compute_importance_map( + ensure_tuple(self.roi_size), mode=mode, sigma_scale=sigma_scale, device=device + ) + if cache_roi_weight_map and self.roi_weight_map is None: + warnings.warn("cache_roi_weight_map=True, but cache is not created. (dynamic roi_size?)") + except BaseException as e: + raise RuntimeError( + "Seems to be OOM. Please try smaller roi_size, or use mode='constant' instead of mode='gaussian'. " + ) from e + def __call__( - self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any - ) -> torch.Tensor: + self, + inputs: torch.Tensor, + network: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor], Dict[Any, torch.Tensor]]], + *args: Any, + **kwargs: Any, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]: """ Args: @@ -178,6 +202,7 @@ def __call__( self.sw_device, self.device, self.progress, + self.roi_weight_map, *args, **kwargs, ) @@ -230,8 +255,13 @@ def __call__(self, inputs: torch.Tensor, network: nn.Module, *args: Any, **kwarg class SliceInferer(SlidingWindowInferer): """ - SliceInferer extends SlidingWindowInferer to provide slice-by-slice (2D) inference - when provided a 3D volume. + SliceInferer extends SlidingWindowInferer to provide slice-by-slice (2D) inference when provided a 3D volume. + A typical use case could be a 2D model (like 2D segmentation UNet) operates on the slices from a 3D volume, + and the output is a 3D volume with 2D slices aggregated. Example:: + + # sliding over the `spatial_dim` + inferer = SliceInferer(roi_size=(64, 256), sw_batch_size=1, spatial_dim=1) + output = inferer(input_volume, net) Args: spatial_dim: Spatial dimension over which the slice-by-slice inference runs on the 3D volume. @@ -248,10 +278,15 @@ class SliceInferer(SlidingWindowInferer): def __init__(self, spatial_dim: int = 0, *args, **kwargs) -> None: self.spatial_dim = spatial_dim super().__init__(*args, **kwargs) + self.orig_roi_size = ensure_tuple(self.roi_size) def __call__( - self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any - ) -> torch.Tensor: + self, + inputs: torch.Tensor, + network: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor], Dict[Any, torch.Tensor]]], + *args: Any, + **kwargs: Any, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]: """ Args: inputs: 3D input for inference @@ -264,21 +299,38 @@ def __call__( # Check if ``roi_size`` tuple is 2D and ``inputs`` tensor is 3D self.roi_size = ensure_tuple(self.roi_size) - if len(self.roi_size) == 2 and len(inputs.shape[2:]) == 3: - self.roi_size = list(self.roi_size) + if len(self.orig_roi_size) == 2 and len(inputs.shape[2:]) == 3: + self.roi_size = list(self.orig_roi_size) self.roi_size.insert(self.spatial_dim, 1) else: - raise RuntimeError("Currently, only 2D `roi_size` with 3D `inputs` tensor is supported.") + raise RuntimeError( + f"Currently, only 2D `roi_size` ({self.orig_roi_size}) with 3D `inputs` tensor (shape={inputs.shape}) is supported." + ) return super().__call__(inputs=inputs, network=lambda x: self.network_wrapper(network, x, *args, **kwargs)) - def network_wrapper(self, network: Callable[..., torch.Tensor], x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + def network_wrapper( + self, + network: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor], Dict[Any, torch.Tensor]]], + x: torch.Tensor, + *args, + **kwargs, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]: """ Wrapper handles inference for 2D models over 3D volume inputs. """ # Pass 4D input [N, C, H, W]/[N, C, D, W]/[N, C, D, H] to the model as it is 2D. x = x.squeeze(dim=self.spatial_dim + 2) out = network(x, *args, **kwargs) + # Unsqueeze the network output so it is [N, C, D, H, W] as expected by # the default SlidingWindowInferer class - return out.unsqueeze(dim=self.spatial_dim + 2) + if isinstance(out, torch.Tensor): + return out.unsqueeze(dim=self.spatial_dim + 2) + + if isinstance(out, Mapping): + for k in out.keys(): + out[k] = out[k].unsqueeze(dim=self.spatial_dim + 2) + return out + + return tuple(out_i.unsqueeze(dim=self.spatial_dim + 2) for out_i in out) diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py index 36e4377bd6..5126b23c0a 100644 --- a/monai/inferers/utils.py +++ b/monai/inferers/utils.py @@ -9,13 +9,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, List, Sequence, Tuple, Union +import warnings +from typing import Any, Callable, Dict, List, Mapping, Sequence, Tuple, Union import torch import torch.nn.functional as F +from monai.data.meta_tensor import MetaTensor from monai.data.utils import compute_importance_map, dense_patch_slices, get_valid_patch_size -from monai.utils import BlendMode, PytorchPadMode, fall_back_tuple, look_up_option, optional_import +from monai.transforms import Resize +from monai.utils import ( + BlendMode, + PytorchPadMode, + convert_data_type, + convert_to_dst_type, + ensure_tuple, + fall_back_tuple, + look_up_option, + optional_import, +) tqdm, _ = optional_import("tqdm", name="tqdm") @@ -26,7 +38,7 @@ def sliding_window_inference( inputs: torch.Tensor, roi_size: Union[Sequence[int], int], sw_batch_size: int, - predictor: Callable[..., torch.Tensor], + predictor: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor], Dict[Any, torch.Tensor]]], overlap: float = 0.25, mode: Union[BlendMode, str] = BlendMode.CONSTANT, sigma_scale: Union[Sequence[float], float] = 0.125, @@ -35,12 +47,21 @@ def sliding_window_inference( sw_device: Union[torch.device, str, None] = None, device: Union[torch.device, str, None] = None, progress: bool = False, + roi_weight_map: Union[torch.Tensor, None] = None, *args: Any, **kwargs: Any, -) -> torch.Tensor: +) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]: """ Sliding window inference on `inputs` with `predictor`. + The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors. + Each output in the tuple or dict value is allowed to have different resolutions with respect to the input. + e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes + could be ([128,64,256], [64,32,128]). + In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still + an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters + so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension). + When roi_size is larger than the inputs' spatial size, the input image are padded during inference. To maintain the same spatial sizes, the output image will be cropped to the original input size. @@ -52,9 +73,16 @@ def sliding_window_inference( corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted to `(32, 64)` if the second spatial dimension size of img is `64`. sw_batch_size: the batch size to run window slices. - predictor: given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)` - should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D]; - where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`. + predictor: given input tensor ``patch_data`` in shape NCHW[D], + The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary + with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D']; + where H'W'[D'] represents the output patch's spatial size, M is the number of output channels, + N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128), + the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)). + In this case, the parameter `overlap` and `roi_size` need to be carefully chosen + to ensure the scaled output ROI sizes are still integers. + If the `predictor`'s input and output spatial sizes are different, + we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension. overlap: Amount of overlap between scans. mode: {``"constant"``, ``"gaussian"``} How to blend output of overlapping windows. Defaults to ``"constant"``. @@ -78,6 +106,8 @@ def sliding_window_inference( set to device=torch.device('cpu') the gpu memory consumption is less and independent of the `inputs` and `roi_size`. Output is on the `device`. progress: whether to print a `tqdm` progress bar. + roi_weight_map: pre-computed (non-negative) weight map for each ROI. + If not given, and ``mode`` is not `constant`, this map will be computed on the fly. args: optional args to be passed to ``predictor``. kwargs: optional keyword args to be passed to ``predictor``. @@ -85,14 +115,14 @@ def sliding_window_inference( - input must be channel-first and have a batch dim, supports N-D sliding window. """ + compute_dtype = inputs.dtype num_spatial_dims = len(inputs.shape) - 2 if overlap < 0 or overlap >= 1: - raise AssertionError("overlap must be >= 0 and < 1.") + raise ValueError("overlap must be >= 0 and < 1.") # determine image spatial size and batch size # Note: all input images must have the same image size and batch size - image_size_ = list(inputs.shape[2:]) - batch_size = inputs.shape[0] + batch_size, _, *image_size_ = inputs.shape if device is None: device = inputs.device @@ -107,7 +137,7 @@ def sliding_window_inference( diff = max(roi_size[k - 2] - inputs.shape[k], 0) half = diff // 2 pad_size.extend([half, diff - half]) - inputs = F.pad(inputs, pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode).value, value=cval) + inputs = F.pad(inputs, pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval) scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap) @@ -117,45 +147,139 @@ def sliding_window_inference( total_slices = num_win * batch_size # total number of windows # Create window-level importance map - importance_map = compute_importance_map( - get_valid_patch_size(image_size, roi_size), mode=mode, sigma_scale=sigma_scale, device=device - ) + valid_patch_size = get_valid_patch_size(image_size, roi_size) + if valid_patch_size == roi_size and (roi_weight_map is not None): + importance_map = roi_weight_map + else: + try: + importance_map = compute_importance_map(valid_patch_size, mode=mode, sigma_scale=sigma_scale, device=device) + except BaseException as e: + raise RuntimeError( + "Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'." + ) from e + importance_map = convert_data_type(importance_map, torch.Tensor, device, compute_dtype)[0] # type: ignore + # handle non-positive weights + min_non_zero = max(importance_map[importance_map != 0].min().item(), 1e-3) + importance_map = torch.clamp(importance_map.to(torch.float32), min=min_non_zero).to(compute_dtype) # Perform predictions - output_image, count_map = torch.tensor(0.0, device=device), torch.tensor(0.0, device=device) - _initialized = False + dict_key, output_image_list, count_map_list = None, [], [] + _initialized_ss = -1 + is_tensor_output = True # whether the predictor's output is a tensor (instead of dict/tuple) + + # for each patch for slice_g in tqdm(range(0, total_slices, sw_batch_size)) if progress else range(0, total_slices, sw_batch_size): slice_range = range(slice_g, min(slice_g + sw_batch_size, total_slices)) unravel_slice = [ [slice(int(idx / num_win), int(idx / num_win) + 1), slice(None)] + list(slices[idx % num_win]) for idx in slice_range ] - window_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device) - seg_prob = predictor(window_data, *args, **kwargs).to(device) # batched patch segmentation - - if not _initialized: # init. buffer at the first iteration - output_classes = seg_prob.shape[1] - output_shape = [batch_size, output_classes] + list(image_size) - # allocate memory to store the full output and the count for overlapping parts - output_image = torch.zeros(output_shape, dtype=torch.float32, device=device) - count_map = torch.zeros(output_shape, dtype=torch.float32, device=device) - _initialized = True - - # store the result in the proper location of the full output. Apply weights from importance map. - for idx, original_idx in zip(slice_range, unravel_slice): - output_image[original_idx] += importance_map * seg_prob[idx - slice_g] - count_map[original_idx] += importance_map + window_data = torch.cat( + [convert_data_type(inputs[win_slice], torch.Tensor)[0] for win_slice in unravel_slice] + ).to(sw_device) + seg_prob_out = predictor(window_data, *args, **kwargs) # batched patch segmentation + + # convert seg_prob_out to tuple seg_prob_tuple, this does not allocate new memory. + seg_prob_tuple: Tuple[torch.Tensor, ...] + if isinstance(seg_prob_out, torch.Tensor): + seg_prob_tuple = (seg_prob_out,) + elif isinstance(seg_prob_out, Mapping): + if dict_key is None: + dict_key = sorted(seg_prob_out.keys()) # track predictor's output keys + seg_prob_tuple = tuple(seg_prob_out[k] for k in dict_key) + is_tensor_output = False + else: + seg_prob_tuple = ensure_tuple(seg_prob_out) + is_tensor_output = False + + # for each output in multi-output list + for ss, seg_prob in enumerate(seg_prob_tuple): + seg_prob = seg_prob.to(device) # BxCxMxNxP or BxCxMxN + + # compute zoom scale: out_roi_size/in_roi_size + zoom_scale = [] + for axis, (img_s_i, out_w_i, in_w_i) in enumerate( + zip(image_size, seg_prob.shape[2:], window_data.shape[2:]) + ): + _scale = out_w_i / float(in_w_i) + if not (img_s_i * _scale).is_integer(): + warnings.warn( + f"For spatial axis: {axis}, output[{ss}] will have non-integer shape. Spatial " + f"zoom_scale between output[{ss}] and input is {_scale}. Please pad inputs." + ) + zoom_scale.append(_scale) + + if _initialized_ss < ss: # init. the ss-th buffer at the first iteration + # construct multi-resolution outputs + output_classes = seg_prob.shape[1] + output_shape = [batch_size, output_classes] + [ + int(image_size_d * zoom_scale_d) for image_size_d, zoom_scale_d in zip(image_size, zoom_scale) + ] + # allocate memory to store the full output and the count for overlapping parts + output_image_list.append(torch.zeros(output_shape, dtype=compute_dtype, device=device)) + count_map_list.append(torch.zeros([1, 1] + output_shape[2:], dtype=compute_dtype, device=device)) + _initialized_ss += 1 + + # resizing the importance_map + resizer = Resize(spatial_size=seg_prob.shape[2:], mode="nearest", anti_aliasing=False) + + # store the result in the proper location of the full output. Apply weights from importance map. + for idx, original_idx in zip(slice_range, unravel_slice): + # zoom roi + original_idx_zoom = list(original_idx) # 4D for 2D image, 5D for 3D image + for axis in range(2, len(original_idx_zoom)): + zoomed_start = original_idx[axis].start * zoom_scale[axis - 2] + zoomed_end = original_idx[axis].stop * zoom_scale[axis - 2] + if not zoomed_start.is_integer() or (not zoomed_end.is_integer()): + warnings.warn( + f"For axis-{axis-2} of output[{ss}], the output roi range is not int. " + f"Input roi range is ({original_idx[axis].start}, {original_idx[axis].stop}). " + f"Spatial zoom_scale between output[{ss}] and input is {zoom_scale[axis - 2]}. " + f"Corresponding output roi range is ({zoomed_start}, {zoomed_end}).\n" + f"Please change overlap ({overlap}) or roi_size ({roi_size[axis-2]}) for axis-{axis-2}. " + "Tips: if overlap*roi_size*zoom_scale is an integer, it usually works." + ) + original_idx_zoom[axis] = slice(int(zoomed_start), int(zoomed_end), None) + importance_map_zoom = resizer(importance_map.unsqueeze(0))[0].to(compute_dtype) + # store results and weights + output_image_list[ss][original_idx_zoom] += importance_map_zoom * seg_prob[idx - slice_g] + count_map_list[ss][original_idx_zoom] += ( + importance_map_zoom.unsqueeze(0).unsqueeze(0).expand(count_map_list[ss][original_idx_zoom].shape) + ) # account for any overlapping sections - output_image = output_image / count_map - - final_slicing: List[slice] = [] - for sp in range(num_spatial_dims): - slice_dim = slice(pad_size[sp * 2], image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2]) - final_slicing.insert(0, slice_dim) - while len(final_slicing) < len(output_image.shape): - final_slicing.insert(0, slice(None)) - return output_image[final_slicing] + for ss in range(len(output_image_list)): + output_image_list[ss] = (output_image_list[ss] / count_map_list.pop(0)).to(compute_dtype) + + # remove padding if image_size smaller than roi_size + for ss, output_i in enumerate(output_image_list): + if torch.isnan(output_i).any() or torch.isinf(output_i).any(): + warnings.warn("Sliding window inference results contain NaN or Inf.") + + zoom_scale = [ + seg_prob_map_shape_d / roi_size_d for seg_prob_map_shape_d, roi_size_d in zip(output_i.shape[2:], roi_size) + ] + + final_slicing: List[slice] = [] + for sp in range(num_spatial_dims): + slice_dim = slice(pad_size[sp * 2], image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2]) + slice_dim = slice( + int(round(slice_dim.start * zoom_scale[num_spatial_dims - sp - 1])), + int(round(slice_dim.stop * zoom_scale[num_spatial_dims - sp - 1])), + ) + final_slicing.insert(0, slice_dim) + while len(final_slicing) < len(output_i.shape): + final_slicing.insert(0, slice(None)) + output_image_list[ss] = output_i[final_slicing] + + if dict_key is not None: # if output of predictor is a dict + final_output = dict(zip(dict_key, output_image_list)) + else: + final_output = tuple(output_image_list) # type: ignore + final_output = final_output[0] if is_tensor_output else final_output # type: ignore + if isinstance(inputs, MetaTensor): + final_output = convert_to_dst_type(final_output, inputs)[0] # type: ignore + return final_output def _get_scan_interval( diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py index 1922996fb6..c3ae941519 100644 --- a/monai/losses/__init__.py +++ b/monai/losses/__init__.py @@ -16,16 +16,20 @@ DiceCELoss, DiceFocalLoss, DiceLoss, + GeneralizedDiceFocalLoss, GeneralizedDiceLoss, GeneralizedWassersteinDiceLoss, MaskedDiceLoss, dice_ce, dice_focal, generalized_dice, + generalized_dice_focal, generalized_wasserstein_dice, ) from .focal_loss import FocalLoss +from .giou_loss import BoxGIoULoss, giou from .image_dissimilarity import GlobalMutualInformationLoss, LocalNormalizedCrossCorrelationLoss from .multi_scale import MultiScaleLoss from .spatial_mask import MaskedLoss from .tversky import TverskyLoss +from .unified_focal_loss import AsymmetricUnifiedFocalLoss diff --git a/monai/losses/dice.py b/monai/losses/dice.py index 610327ef63..892d71d06a 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -291,9 +291,9 @@ def __init__( self.batch = batch def w_func(self, grnd): - if self.w_type == Weight.SIMPLE: + if self.w_type == str(Weight.SIMPLE): return torch.reciprocal(grnd) - if self.w_type == Weight.SQUARE: + if self.w_type == str(Weight.SQUARE): return torch.reciprocal(grnd * grnd) return torch.ones_like(grnd) @@ -796,8 +796,6 @@ def __init__( """ super().__init__() self.dice = DiceLoss( - include_background=include_background, - to_onehot_y=to_onehot_y, sigmoid=sigmoid, softmax=softmax, other_act=other_act, @@ -808,19 +806,15 @@ def __init__( smooth_dr=smooth_dr, batch=batch, ) - self.focal = FocalLoss( - include_background=include_background, - to_onehot_y=to_onehot_y, - gamma=gamma, - weight=focal_weight, - reduction=reduction, - ) + self.focal = FocalLoss(gamma=gamma, weight=focal_weight, reduction=reduction) if lambda_dice < 0.0: raise ValueError("lambda_dice should be no less than 0.0.") if lambda_focal < 0.0: raise ValueError("lambda_focal should be no less than 0.0.") self.lambda_dice = lambda_dice self.lambda_focal = lambda_focal + self.to_onehot_y = to_onehot_y + self.include_background = include_background def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """ @@ -837,14 +831,137 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if len(input.shape) != len(target.shape): raise ValueError("the number of dimensions for input and target should be the same.") + n_pred_ch = input.shape[1] + + if self.to_onehot_y: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + 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.") + else: + # if skipping background, removing first channel + target = target[:, 1:] + input = input[:, 1:] + dice_loss = self.dice(input, target) focal_loss = self.focal(input, target) total_loss: torch.Tensor = self.lambda_dice * dice_loss + self.lambda_focal * focal_loss return total_loss +class GeneralizedDiceFocalLoss(torch.nn.modules.loss._Loss): + """Compute both Generalized Dice Loss and Focal Loss, and return their weighted average. The details of Generalized Dice Loss + and Focal Loss are available at ``monai.losses.GeneralizedDiceLoss`` and ``monai.losses.FocalLoss``. + + Args: + include_background (bool, optional): if False channel index 0 (background category) is excluded from the calculation. + Defaults to True. + to_onehot_y (bool, optional): whether to convert `y` into the one-hot format. Defaults to False. + sigmoid (bool, optional): if True, apply a sigmoid function to the prediction. Defaults to False. + softmax (bool, optional): if True, apply a softmax function to the prediction. Defaults to False. + other_act (Optional[Callable], optional): if don't want to use sigmoid or softmax, use other callable + function to execute other activation layers. Defaults to None. + w_type (Union[Weight, str], optional): {``"square"``, ``"simple"``, ``"uniform"``}. Type of function to transform + ground-truth volume to a weight factor. Defaults to ``"square"``. + reduction (Union[LossReduction, str], optional): {``"none"``, ``"mean"``, ``"sum"``}. Specified 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. + smooth_nr (float, optional): a small constant added to the numerator to avoid zero. Defaults to 1e-5. + smooth_dr (float, optional): a small constant added to the denominator to avoid nan. Defaults to 1e-5. + batch (bool, optional): whether to sum the intersection and union areas over the batch dimension before the dividing. + Defaults to False, i.e., the areas are computed for each item in the batch. + gamma (float, optional): value of the exponent gamma in the definition of the Focal loss. Defaults to 2.0. + focal_weight (Optional[Union[Sequence[float], float, int, torch.Tensor]], optional): weights to apply to + the voxels of each class. If None no weights are applied. The input can be a single value + (same weight for all classes), a sequence of values (the length of the sequence hould be the same as + the number of classes). Defaults to None. + lambda_gdl (float, optional): the trade-off weight value for Generalized Dice Loss. The value should be + no less than 0.0. Defaults to 1.0. + lambda_focal (float, optional): the trade-off weight value for Focal Loss. The value should be no less + than 0.0. Defaults to 1.0. + + Raises: + ValueError: if either `lambda_gdl` or `lambda_focal` is less than 0. + """ + + def __init__( + self, + include_background: bool = True, + to_onehot_y: bool = False, + sigmoid: bool = False, + softmax: bool = False, + other_act: Optional[Callable] = None, + w_type: Union[Weight, str] = Weight.SQUARE, + reduction: Union[LossReduction, str] = LossReduction.MEAN, + smooth_nr: float = 1e-5, + smooth_dr: float = 1e-5, + batch: bool = False, + gamma: float = 2.0, + focal_weight: Optional[Union[Sequence[float], float, int, torch.Tensor]] = None, + lambda_gdl: float = 1.0, + lambda_focal: float = 1.0, + ) -> None: + super().__init__() + self.generalized_dice = GeneralizedDiceLoss( + include_background=include_background, + to_onehot_y=to_onehot_y, + sigmoid=sigmoid, + softmax=softmax, + other_act=other_act, + w_type=w_type, + reduction=reduction, + smooth_nr=smooth_nr, + smooth_dr=smooth_dr, + batch=batch, + ) + self.focal = FocalLoss( + include_background=include_background, + to_onehot_y=to_onehot_y, + gamma=gamma, + weight=focal_weight, + reduction=reduction, + ) + if lambda_gdl < 0.0: + raise ValueError("lambda_gdl should be no less than 0.0.") + if lambda_focal < 0.0: + raise ValueError("lambda_focal should be no less than 0.0.") + self.lambda_gdl = lambda_gdl + self.lambda_focal = lambda_focal + + def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """ + Args: + input (torch.Tensor): the shape should be BNH[WD]. The input should be the original logits + due to the restriction of ``monai.losses.FocalLoss``. + target (torch.Tensor): the shape should be BNH[WD] or B1H[WD]. + + Raises: + ValueError: When the input and target tensors have different numbers of dimensions, or the target + channel isn't either one-hot encoded or categorical with the same shape of the input. + + Returns: + torch.Tensor: value of the loss. + """ + if input.dim() != target.dim(): + raise ValueError( + f"Input - {input.shape} - and target - {target.shape} - must have the same number of dimensions." + ) + + gdl_loss = self.generalized_dice(input, target) + focal_loss = self.focal(input, target) + total_loss: torch.Tensor = self.lambda_gdl * gdl_loss + self.lambda_focal * focal_loss + return total_loss + + Dice = DiceLoss dice_ce = DiceCELoss dice_focal = DiceFocalLoss generalized_dice = GeneralizedDiceLoss +generalized_dice_focal = GeneralizedDiceFocalLoss generalized_wasserstein_dice = GeneralizedWassersteinDiceLoss diff --git a/monai/losses/giou_loss.py b/monai/losses/giou_loss.py new file mode 100644 index 0000000000..ec7e358f42 --- /dev/null +++ b/monai/losses/giou_loss.py @@ -0,0 +1,68 @@ +# 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 typing import Union + +import torch +from torch.nn.modules.loss import _Loss + +from monai.data.box_utils import COMPUTE_DTYPE, box_pair_giou +from monai.utils import LossReduction + + +class BoxGIoULoss(_Loss): + + """ + Compute the generalized intersection over union (GIoU) loss of a pair of boxes. + The two inputs should have the same shape. giou_loss = 1.0 - giou + + The range of GIoU is (-1.0, 1.0]. Thus the range of GIoU loss is [0.0, 2.0). + + Args: + 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. + """ + + def __init__(self, reduction: Union[LossReduction, str] = LossReduction.MEAN) -> None: + super().__init__(reduction=LossReduction(reduction).value) + + def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """ + Args: + input: predicted bounding boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + target: GT bounding boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + + Raises: + ValueError: When the two inputs have different shape. + """ + if target.shape != input.shape: + raise ValueError(f"ground truth has different shape ({target.shape}) from input ({input.shape})") + + box_dtype = input.dtype + giou: torch.Tensor = box_pair_giou( # type: ignore + target.to(dtype=COMPUTE_DTYPE), input.to(dtype=COMPUTE_DTYPE) + ) + loss: torch.Tensor = 1.0 - giou + if self.reduction == LossReduction.MEAN.value: + loss = loss.mean() + elif self.reduction == LossReduction.SUM.value: + loss = loss.sum() + elif self.reduction == LossReduction.NONE.value: + pass + else: + raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') + return loss.to(box_dtype) + + +giou = BoxGIoULoss diff --git a/monai/losses/unified_focal_loss.py b/monai/losses/unified_focal_loss.py new file mode 100644 index 0000000000..1f6d51bead --- /dev/null +++ b/monai/losses/unified_focal_loss.py @@ -0,0 +1,239 @@ +# 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. + +import warnings +from typing import Union + +import torch +from torch.nn.modules.loss import _Loss + +from monai.networks import one_hot +from monai.utils import LossReduction + + +class AsymmetricFocalTverskyLoss(_Loss): + """ + AsymmetricFocalTverskyLoss is a variant of FocalTverskyLoss, which attentions to the foreground class. + + Actually, it's only supported for binary image segmentation now. + + Reimplementation of the Asymmetric Focal Tversky Loss described in: + + - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation", + Michael Yeung, Computerized Medical Imaging and Graphics + """ + + def __init__( + self, + to_onehot_y: bool = False, + delta: float = 0.7, + gamma: float = 0.75, + epsilon: float = 1e-7, + reduction: Union[LossReduction, str] = LossReduction.MEAN, + ) -> None: + """ + Args: + to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False. + delta : weight of the background. Defaults to 0.7. + gamma : value of the exponent gamma in the definition of the Focal loss . Defaults to 0.75. + epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7. + """ + super().__init__(reduction=LossReduction(reduction).value) + self.to_onehot_y = to_onehot_y + self.delta = delta + self.gamma = gamma + self.epsilon = epsilon + + 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.") + else: + y_true = one_hot(y_true, num_classes=n_pred_ch) + + if y_true.shape != y_pred.shape: + raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})") + + # clip the prediction to avoid NaN + y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon) + axis = list(range(2, len(y_pred.shape))) + + # Calculate true positives (tp), false negatives (fn) and false positives (fp) + tp = torch.sum(y_true * y_pred, dim=axis) + fn = torch.sum(y_true * (1 - y_pred), dim=axis) + fp = torch.sum((1 - y_true) * y_pred, dim=axis) + dice_class = (tp + self.epsilon) / (tp + self.delta * fn + (1 - self.delta) * fp + self.epsilon) + + # Calculate losses separately for each class, enhancing both classes + back_dice = 1 - dice_class[:, 0] + fore_dice = (1 - dice_class[:, 1]) * torch.pow(1 - dice_class[:, 1], -self.gamma) + + # Average class scores + loss = torch.mean(torch.stack([back_dice, fore_dice], dim=-1)) + return loss + + +class AsymmetricFocalLoss(_Loss): + """ + AsymmetricFocalLoss is a variant of FocalTverskyLoss, which attentions to the foreground class. + + Actually, it's only supported for binary image segmentation now. + + Reimplementation of the Asymmetric Focal Loss described in: + + - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation", + Michael Yeung, Computerized Medical Imaging and Graphics + """ + + def __init__( + self, + to_onehot_y: bool = False, + delta: float = 0.7, + gamma: float = 2, + epsilon: float = 1e-7, + reduction: Union[LossReduction, str] = LossReduction.MEAN, + ): + """ + Args: + to_onehot_y : whether to convert `y` into the one-hot format. Defaults to False. + delta : weight of the background. Defaults to 0.7. + gamma : value of the exponent gamma in the definition of the Focal loss . Defaults to 0.75. + epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7. + """ + super().__init__(reduction=LossReduction(reduction).value) + self.to_onehot_y = to_onehot_y + self.delta = delta + self.gamma = gamma + self.epsilon = epsilon + + 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.") + else: + y_true = one_hot(y_true, num_classes=n_pred_ch) + + if y_true.shape != y_pred.shape: + raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})") + + y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon) + cross_entropy = -y_true * torch.log(y_pred) + + back_ce = torch.pow(1 - y_pred[:, 0], self.gamma) * cross_entropy[:, 0] + back_ce = (1 - self.delta) * back_ce + + fore_ce = cross_entropy[:, 1] + fore_ce = self.delta * fore_ce + + loss = torch.mean(torch.sum(torch.stack([back_ce, fore_ce], dim=1), dim=1)) + return loss + + +class AsymmetricUnifiedFocalLoss(_Loss): + """ + AsymmetricUnifiedFocalLoss is a variant of Focal Loss. + + Actually, it's only supported for binary image segmentation now + + Reimplementation of the Asymmetric Unified Focal Tversky Loss described in: + + - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation", + Michael Yeung, Computerized Medical Imaging and Graphics + """ + + def __init__( + self, + to_onehot_y: bool = False, + num_classes: int = 2, + weight: float = 0.5, + gamma: float = 0.5, + delta: float = 0.7, + reduction: Union[LossReduction, str] = LossReduction.MEAN, + ): + """ + Args: + to_onehot_y : whether to convert `y` into the one-hot format. Defaults to False. + num_classes : number of classes, it only supports 2 now. Defaults to 2. + delta : weight of the background. Defaults to 0.7. + gamma : value of the exponent gamma in the definition of the Focal loss. Defaults to 0.75. + epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7. + weight : weight for each loss function, if it's none it's 0.5. Defaults to None. + + Example: + >>> import torch + >>> from monai.losses import AsymmetricUnifiedFocalLoss + >>> pred = torch.ones((1,1,32,32), dtype=torch.float32) + >>> grnd = torch.ones((1,1,32,32), dtype=torch.int64) + >>> fl = AsymmetricUnifiedFocalLoss(to_onehot_y=True) + >>> fl(pred, grnd) + """ + super().__init__(reduction=LossReduction(reduction).value) + self.to_onehot_y = to_onehot_y + self.num_classes = num_classes + self.gamma = gamma + self.delta = delta + self.weight: float = weight + self.asy_focal_loss = AsymmetricFocalLoss(gamma=self.gamma, delta=self.delta) + self.asy_focal_tversky_loss = AsymmetricFocalTverskyLoss(gamma=self.gamma, delta=self.delta) + + # TODO: Implement this function to support multiple classes segmentation + def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor: + """ + Args: + y_pred : the shape should be BNH[WD], where N is the number of classes. + It only supports binary segmentation. + The input should be the original logits since it will be transformed by + a sigmoid in the forward function. + y_true : the shape should be BNH[WD], where N is the number of classes. + It only supports binary segmentation. + + Raises: + ValueError: When input and target are different shape + ValueError: When len(y_pred.shape) != 4 and len(y_pred.shape) != 5 + ValueError: When num_classes + ValueError: When the number of classes entered does not match the expected number + """ + if y_pred.shape != y_true.shape: + raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})") + + if len(y_pred.shape) != 4 and len(y_pred.shape) != 5: + raise ValueError(f"input shape must be 4 or 5, but got {y_pred.shape}") + + if y_pred.shape[1] == 1: + y_pred = one_hot(y_pred, num_classes=self.num_classes) + y_true = one_hot(y_true, num_classes=self.num_classes) + + if torch.max(y_true) != self.num_classes - 1: + raise ValueError(f"Pelase make sure the number of classes is {self.num_classes-1}") + + 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.") + else: + y_true = one_hot(y_true, num_classes=n_pred_ch) + + asy_focal_loss = self.asy_focal_loss(y_pred, y_true) + asy_focal_tversky_loss = self.asy_focal_tversky_loss(y_pred, y_true) + + loss: torch.Tensor = self.weight * asy_focal_loss + (1 - self.weight) * asy_focal_tversky_loss + + if self.reduction == LossReduction.SUM.value: + return torch.sum(loss) # sum over the batch and channel dims + if self.reduction == LossReduction.NONE.value: + return loss # returns [N, num_classes] losses + if self.reduction == LossReduction.MEAN.value: + return torch.mean(loss) + raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py index d18c20f7b2..750f0f3552 100644 --- a/monai/metrics/__init__.py +++ b/monai/metrics/__init__.py @@ -12,10 +12,12 @@ from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix from .cumulative_average import CumulativeAverage from .froc import compute_fp_tp_probs, compute_froc_curve_data, compute_froc_score +from .generalized_dice import GeneralizedDiceScore, compute_generalized_dice from .hausdorff_distance import HausdorffDistanceMetric, compute_hausdorff_distance, compute_percent_hausdorff_distance from .meandice import DiceMetric, compute_meandice from .metric import Cumulative, CumulativeIterationMetric, IterationMetric, Metric from .regression import MAEMetric, MSEMetric, PSNRMetric, RMSEMetric from .rocauc import ROCAUCMetric, compute_roc_auc +from .surface_dice import SurfaceDiceMetric, compute_surface_dice from .surface_distance import SurfaceDistanceMetric, compute_average_surface_distance -from .utils import do_metric_reduction, get_mask_edges, get_surface_distance, ignore_background +from .utils import do_metric_reduction, get_mask_edges, get_surface_distance, ignore_background, is_binary_tensor diff --git a/monai/metrics/confusion_matrix.py b/monai/metrics/confusion_matrix.py index 320f657537..bb195b0cb8 100644 --- a/monai/metrics/confusion_matrix.py +++ b/monai/metrics/confusion_matrix.py @@ -14,7 +14,7 @@ import torch -from monai.metrics.utils import do_metric_reduction, ignore_background +from monai.metrics.utils import do_metric_reduction, ignore_background, is_binary_tensor from monai.utils import MetricReduction, ensure_tuple from .metric import CumulativeIterationMetric @@ -30,7 +30,7 @@ class ConfusionMatrixMetric(CumulativeIterationMetric): The `include_background` parameter can be set to ``False`` for an instance to exclude the first category (channel index 0) which is by convention assumed to be background. 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. + background. Args: include_background: whether to skip metric computation on the first channel of @@ -47,7 +47,7 @@ class ConfusionMatrixMetric(CumulativeIterationMetric): returned with the same order as input names when calling the class. compute_sample: when reducing, if ``True``, each sample's metric will be computed based on each confusion matrix first. if ``False``, compute reduction on the confusion matrices first, defaults to ``False``. - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns [(metric, not_nans), ...]. If False, @@ -84,13 +84,9 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor ValueError: when `y` is not a binarized tensor. ValueError: when `y_pred` has less than two dimensions. """ - if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): - raise ValueError("y_pred and y must be PyTorch Tensor.") - # check binarized input - if not torch.all(y_pred.byte() == y_pred): - warnings.warn("y_pred should be a binarized tensor.") - if not torch.all(y.byte() == y): - raise ValueError("y should be a binarized tensor.") + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + # check dimension dims = y_pred.ndimension() if dims < 2: @@ -102,10 +98,19 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor return get_confusion_matrix(y_pred=y_pred, y=y, include_background=self.include_background) - def aggregate(self): + def aggregate( # type: ignore + self, compute_sample: bool = False, reduction: Union[MetricReduction, str, None] = None + ): """ Execute reduction for the confusion matrix values. + Args: + compute_sample: when reducing, if ``True``, each sample's metric will be computed based on each confusion matrix first. + if ``False``, compute reduction on the confusion matrices first, defaults to ``False``. + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + """ data = self.get_buffer() if not isinstance(data, torch.Tensor): @@ -113,11 +118,11 @@ def aggregate(self): results = [] for metric_name in self.metric_name: - if self.compute_sample: + if compute_sample or self.compute_sample: sub_confusion_matrix = compute_confusion_matrix_metric(metric_name, data) - f, not_nans = do_metric_reduction(sub_confusion_matrix, self.reduction) + f, not_nans = do_metric_reduction(sub_confusion_matrix, reduction or self.reduction) else: - f, not_nans = do_metric_reduction(data, self.reduction) + f, not_nans = do_metric_reduction(data, reduction or self.reduction) f = compute_confusion_matrix_metric(metric_name, f) if self.get_not_nans: results.append((f, not_nans)) @@ -313,7 +318,7 @@ def check_confusion_matrix_metric_name(metric_name: str): return "mcc" if metric_name in ["fowlkes_mallows_index", "fm"]: return "fm" - if metric_name in ["informedness", "bookmaker_informedness", "bm"]: + if metric_name in ["informedness", "bookmaker_informedness", "bm", "youden_index", "youden"]: return "bm" if metric_name in ["markedness", "deltap", "mk"]: return "mk" diff --git a/monai/metrics/generalized_dice.py b/monai/metrics/generalized_dice.py new file mode 100644 index 0000000000..0e1f61ac68 --- /dev/null +++ b/monai/metrics/generalized_dice.py @@ -0,0 +1,183 @@ +# 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 typing import Union + +import torch + +from monai.metrics.utils import do_metric_reduction, ignore_background +from monai.utils import MetricReduction, Weight, look_up_option + +from .metric import CumulativeIterationMetric +from .utils import is_binary_tensor + + +class GeneralizedDiceScore(CumulativeIterationMetric): + """Compute the Generalized Dice Score metric between tensors, as the complement of the Generalized Dice Loss defined in: + + Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning + loss function for highly unbalanced segmentations. DLMIA 2017. + + The inputs `y_pred` and `y` are expected to be one-hot, binarized channel-first + or batch-first tensors, i.e., CHW[D] or BCHW[D]. + + Args: + include_background (bool, optional): whether to include the background class (assumed to be in channel 0), in the + score computation. Defaults to True. + reduction (str, optional): define mode of reduction to the metrics. Available reduction modes: + {``"none"``, ``"mean_batch"``, ``"sum_batch"``}. Default to ``"mean_batch"``. If "none", will not do reduction. + weight_type (Union[Weight, str], optional): {``"square"``, ``"simple"``, ``"uniform"``}. Type of function to transform + ground truth volume into a weight factor. Defaults to ``"square"``. + + Raises: + ValueError: when the `weight_type` is not one of {``"none"``, ``"mean"``, ``"sum"``}. + """ + + def __init__( + self, + include_background: bool = True, + reduction: Union[MetricReduction, str] = MetricReduction.MEAN_BATCH, + weight_type: Union[Weight, str] = Weight.SQUARE, + ) -> None: + super().__init__() + self.include_background = include_background + reduction_options = [ + "none", + "mean_batch", + "sum_batch", + MetricReduction.NONE, + MetricReduction.MEAN_BATCH, + MetricReduction.SUM_BATCH, + ] + self.reduction = reduction + if self.reduction not in reduction_options: + raise ValueError(f"reduction must be one of {reduction_options}") + self.weight_type = look_up_option(weight_type, Weight) + + def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignore + """Computes the Generalized Dice Score and returns a tensor with its per image values. + + Args: + y_pred (torch.Tensor): binarized segmentation model output. It must be in one-hot format and in the NCHW[D] format, + where N is the batch dimension, C is the channel dimension, and the remaining are the spatial dimensions. + y (torch.Tensor): binarized ground-truth. It must be in one-hot format and have the same shape as `y_pred`. + + Raises: + ValueError: if `y_pred` or `y` is not a binarized PyTorch tensor, if `y_pred` and `y` have less than + three dimensions, or `y_pred` and `y` don't have the same shape. + """ + return compute_generalized_dice( + y_pred=y_pred, y=y, include_background=self.include_background, weight_type=self.weight_type + ) + + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore + """ + Execute reduction logic for the output of `compute_generalized_dice`. + + Args: + reduction (Union[MetricReduction, str, None], optional): define mode of reduction to the metrics. + Available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``}. + Defaults to ``"mean"``. If "none", will not do reduction. + """ + data = self.get_buffer() + if not isinstance(data, torch.Tensor): + raise ValueError("The data to aggregate must be a PyTorch Tensor.") + + # Validate reduction argument if specified + if reduction is not None: + reduction_options = ["none", "mean", "sum", "mean_batch", "sum_batch"] + if reduction not in reduction_options: + raise ValueError(f"reduction must be one of {reduction_options}") + + # Do metric reduction and return + f, _ = do_metric_reduction(data, reduction or self.reduction) + + return f + + +def compute_generalized_dice( + y_pred: torch.Tensor, + y: torch.Tensor, + include_background: bool = True, + weight_type: Union[Weight, str] = Weight.SQUARE, +) -> torch.Tensor: + """Computes the Generalized Dice Score and returns a tensor with its per image values. + + Args: + y_pred (torch.Tensor): binarized segmentation model output. It should be binarized, in one-hot format + and in the NCHW[D] format, where N is the batch dimension, C is the channel dimension, and the + remaining are the spatial dimensions. + y (torch.Tensor): binarized ground-truth. It should be binarized, in one-hot format and have the same shape as `y_pred`. + include_background (bool, optional): whether to skip score computation on the first channel of the + predicted output. Defaults to True. + weight_type (Union[Weight, str], optional): {``"square"``, ``"simple"``, ``"uniform"``}. Type of function to + transform ground truth volume into a weight factor. Defaults to ``"square"``. + + Returns: + torch.Tensor: per batch and per class Generalized Dice Score, i.e., with the shape [batch_size, num_classes]. + + Raises: + ValueError: if `y_pred` or `y` are not PyTorch tensors, if `y_pred` and `y` have less than three dimensions, + or `y_pred` and `y` don't have the same shape. + """ + # Ensure tensors are binarized + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + + # Ensure tensors have at least 3 dimensions and have the same shape + dims = y_pred.dim() + if dims < 3: + raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.") + if y.shape != y_pred.shape: + raise ValueError(f"y_pred - {y_pred.shape} - and y - {y.shape} - should have the same shapes.") + + # Ignore background, if needed + if not include_background: + y_pred, y = ignore_background(y_pred=y_pred, y=y) + + # Reducing only spatial dimensions (not batch nor channels), compute the intersection and non-weighted denominator + reduce_axis = list(range(2, y_pred.dim())) + intersection = torch.sum(y * y_pred, dim=reduce_axis) + y_o = torch.sum(y, dim=reduce_axis) + y_pred_o = torch.sum(y_pred, dim=reduce_axis) + denominator = y_o + y_pred_o + + # Set the class weights + weight_type = look_up_option(weight_type, Weight) + if weight_type == Weight.SIMPLE: + w = torch.reciprocal(y_o.float()) + elif weight_type == Weight.SQUARE: + w = torch.reciprocal(y_o.float() * y_o.float()) + else: + w = torch.ones_like(y_o.float()) + + # Replace infinite values for non-appearing classes by the maximum weight + for b in w: + infs = torch.isinf(b) + b[infs] = 0 + b[infs] = torch.max(b) + + # Compute the weighted numerator and denominator, summing along the class axis + numer = 2.0 * (intersection * w).sum(dim=1) + denom = (denominator * w).sum(dim=1) + + # Compute the score + generalized_dice_score = numer / denom + + # Handle zero deivision. Where denom == 0 and the prediction volume is 0, score is 1. + # Where denom == 0 but the prediction volume is not 0, score is 0 + y_pred_o = y_pred_o.sum(dim=-1) + denom_zeros = denom == 0 + generalized_dice_score[denom_zeros] = torch.where( + (y_pred_o == 0)[denom_zeros], torch.tensor(1.0), torch.tensor(0.0) + ) + + return generalized_dice_score diff --git a/monai/metrics/hausdorff_distance.py b/monai/metrics/hausdorff_distance.py index 5ce739d1f4..1caec2d919 100644 --- a/monai/metrics/hausdorff_distance.py +++ b/monai/metrics/hausdorff_distance.py @@ -15,7 +15,13 @@ import numpy as np import torch -from monai.metrics.utils import do_metric_reduction, get_mask_edges, get_surface_distance, ignore_background +from monai.metrics.utils import ( + do_metric_reduction, + get_mask_edges, + get_surface_distance, + ignore_background, + is_binary_tensor, +) from monai.utils import MetricReduction from .metric import CumulativeIterationMetric @@ -42,7 +48,7 @@ class HausdorffDistanceMetric(CumulativeIterationMetric): percentile of the Hausdorff Distance rather than the maximum result will be achieved. Defaults to ``None``. directed: whether to calculate directed Hausdorff distance. Defaults to ``False``. - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans). @@ -80,12 +86,9 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor ValueError: when `y` is not a binarized tensor. ValueError: when `y_pred` has less than three dimensions. """ - if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): - raise ValueError("y_pred and y must be PyTorch Tensor.") - if not torch.all(y_pred.byte() == y_pred): - warnings.warn("y_pred should be a binarized tensor.") - if not torch.all(y.byte() == y): - warnings.warn("y should be a binarized tensor.") + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + dims = y_pred.ndimension() if dims < 3: raise ValueError("y_pred should have at least three dimensions.") @@ -99,17 +102,22 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor directed=self.directed, ) - def aggregate(self): + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore """ Execute reduction logic for the output of `compute_hausdorff_distance`. + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + """ data = self.get_buffer() if not isinstance(data, torch.Tensor): raise ValueError("the data to aggregate must be PyTorch Tensor.") # do metric reduction - f, not_nans = do_metric_reduction(data, self.reduction) + f, not_nans = do_metric_reduction(data, reduction or self.reduction) return (f, not_nans) if self.get_not_nans else f diff --git a/monai/metrics/meandice.py b/monai/metrics/meandice.py index 4179420804..c86efa3c52 100644 --- a/monai/metrics/meandice.py +++ b/monai/metrics/meandice.py @@ -9,12 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import warnings from typing import Union import torch -from monai.metrics.utils import do_metric_reduction, ignore_background +from monai.metrics.utils import do_metric_reduction, ignore_background, is_binary_tensor from monai.utils import MetricReduction from .metric import CumulativeIterationMetric @@ -22,24 +21,27 @@ class DiceMetric(CumulativeIterationMetric): """ - Compute average Dice loss between two tensors. It can support both multi-classes and multi-labels tasks. + Compute average Dice score between two tensors. It can support both multi-classes and multi-labels tasks. Input `y_pred` is compared with ground truth `y`. `y_preds` is expected to have binarized predictions and `y` should be in one-hot format. You can use suitable transforms in ``monai.transforms.post`` first to achieve binarized values. - The `include_background` parameter can be set to ``False`` for an instance of DiceLoss to exclude + The `include_background` parameter can be set to ``False`` to exclude the first category (channel index 0) which is by convention assumed to be background. 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. + background. `y_preds` and `y` can be a list of channel-first Tensor (CHW[D]) or a batch-first Tensor (BCHW[D]). Args: include_background: whether to skip Dice computation on the first channel of the predicted output. Defaults to ``True``. - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans). Here `not_nans` count the number of not nans for the metric, thus its shape equals to the shape of the metric. + ignore_empty: whether to ignore empty ground truth cases during calculation. + If `True`, NaN value will be set for empty ground truth cases. + If `False`, 1 will be set if the predictions of empty ground truth cases are also empty. """ @@ -48,11 +50,13 @@ def __init__( include_background: bool = True, reduction: Union[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): # type: ignore """ @@ -67,33 +71,39 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor ValueError: when `y` is not a binarized tensor. ValueError: when `y_pred` has less than three dimensions. """ - if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): - raise ValueError("y_pred and y must be PyTorch Tensor.") - if not torch.all(y_pred.byte() == y_pred): - warnings.warn("y_pred should be a binarized tensor.") - if not torch.all(y.byte() == y): - warnings.warn("y should be a binarized tensor.") + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + dims = y_pred.ndimension() if dims < 3: - raise ValueError("y_pred should have at least three dimensions.") + raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.") # compute dice (BxC) for each channel for each batch - return compute_meandice(y_pred=y_pred, y=y, include_background=self.include_background) + return compute_meandice( + y_pred=y_pred, y=y, include_background=self.include_background, ignore_empty=self.ignore_empty + ) - def aggregate(self): + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore """ Execute reduction logic for the output of `compute_meandice`. + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + """ data = self.get_buffer() if not isinstance(data, torch.Tensor): raise ValueError("the data to aggregate must be PyTorch Tensor.") # do metric reduction - f, not_nans = do_metric_reduction(data, self.reduction) + f, not_nans = do_metric_reduction(data, reduction or self.reduction) return (f, not_nans) if self.get_not_nans else f -def compute_meandice(y_pred: torch.Tensor, y: torch.Tensor, include_background: bool = True) -> torch.Tensor: +def compute_meandice( + y_pred: torch.Tensor, y: torch.Tensor, include_background: bool = True, ignore_empty: bool = True +) -> torch.Tensor: """Computes Dice score metric from full size Tensor and collects average. Args: @@ -104,6 +114,9 @@ def compute_meandice(y_pred: torch.Tensor, y: torch.Tensor, include_background: The values should be binarized. include_background: whether to skip Dice computation on the first channel of the predicted output. Defaults to True. + ignore_empty: whether to ignore empty ground truth cases during calculation. + If `True`, NaN value will be set for empty ground truth cases. + If `False`, 1 will be set if the predictions of empty ground truth cases are also empty. Returns: Dice scores per batch and per class, (shape [batch_size, num_classes]). @@ -131,4 +144,6 @@ def compute_meandice(y_pred: torch.Tensor, y: torch.Tensor, include_background: y_pred_o = torch.sum(y_pred, dim=reduce_axis) denominator = y_o + y_pred_o - return torch.where(y_o > 0, (2.0 * intersection) / denominator, torch.tensor(float("nan"), device=y_o.device)) + if ignore_empty is True: + return torch.where(y_o > 0, (2.0 * intersection) / denominator, torch.tensor(float("nan"), device=y_o.device)) + return torch.where(denominator > 0, (2.0 * intersection) / denominator, torch.tensor(1.0, device=y_o.device)) diff --git a/monai/metrics/metric.py b/monai/metrics/metric.py index 7782c4c468..fa8b3354de 100644 --- a/monai/metrics/metric.py +++ b/monai/metrics/metric.py @@ -274,7 +274,10 @@ def get_buffer(self): """ self._sync() - return self._synced_tensors[0] if len(self._synced_tensors) == 1 else self._synced_tensors + if self._synced_tensors is None: + return self._synced_tensors + buffers = [x.detach().clone() if isinstance(x, torch.Tensor) else x for x in self._synced_tensors] + return buffers[0] if len(buffers) == 1 else buffers class CumulativeIterationMetric(Cumulative, IterationMetric): diff --git a/monai/metrics/regression.py b/monai/metrics/regression.py index d5733eee97..62f5fa939e 100644 --- a/monai/metrics/regression.py +++ b/monai/metrics/regression.py @@ -30,7 +30,7 @@ class RegressionMetric(CumulativeIterationMetric): `y_preds` and `y` can be a list of channel-first Tensor (CHW[D]) or a batch-first Tensor (BCHW[D]). Args: - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans). @@ -45,12 +45,18 @@ def __init__( self.reduction = reduction self.get_not_nans = get_not_nans - def aggregate(self): + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore + """ + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + """ data = self.get_buffer() if not isinstance(data, torch.Tensor): raise ValueError("the data to aggregate must be PyTorch Tensor.") - f, not_nans = do_metric_reduction(data, self.reduction) + f, not_nans = do_metric_reduction(data, reduction or self.reduction) return (f, not_nans) if self.get_not_nans else f def _check_shape(self, y_pred: torch.Tensor, y: torch.Tensor) -> None: diff --git a/monai/metrics/rocauc.py b/monai/metrics/rocauc.py index 221fc50272..bd3cdb1203 100644 --- a/monai/metrics/rocauc.py +++ b/monai/metrics/rocauc.py @@ -49,10 +49,14 @@ def __init__(self, average: Union[Average, str] = Average.MACRO) -> None: def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignore return y_pred, y - def aggregate(self): + def aggregate(self, average: Union[Average, str, None] = None): # type: ignore """ - As AUC metric needs to execute on the overall data, so usually users accumulate `y_pred` and `y` - of every iteration, then execute real computation and reduction on the accumulated data. + Typically `y_pred` and `y` are stored in the cumulative buffers at each iteration, + This function reads the buffers and computes the area under the ROC. + + Args: + average: {``"macro"``, ``"weighted"``, ``"micro"``, ``"none"``} + Type of averaging performed if not binary classification. Defaults to `self.average`. """ y_pred, y = self.get_buffer() @@ -60,7 +64,7 @@ def aggregate(self): if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): raise ValueError("y_pred and y must be PyTorch Tensor.") - return compute_roc_auc(y_pred=y_pred, y=y, average=self.average) + return compute_roc_auc(y_pred=y_pred, y=y, average=average or self.average) def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float: diff --git a/monai/metrics/surface_dice.py b/monai/metrics/surface_dice.py new file mode 100644 index 0000000000..f964b0472a --- /dev/null +++ b/monai/metrics/surface_dice.py @@ -0,0 +1,238 @@ +# 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. + +import warnings +from typing import List, Union + +import numpy as np +import torch + +from monai.metrics.utils import do_metric_reduction, get_mask_edges, get_surface_distance, ignore_background +from monai.utils import MetricReduction, convert_data_type + +from .metric import CumulativeIterationMetric + + +class SurfaceDiceMetric(CumulativeIterationMetric): + """ + Computes the Normalized Surface Distance (NSD) for each batch sample and class of + predicted segmentations `y_pred` and corresponding reference segmentations `y` according to equation :eq:`nsd`. + This implementation supports 2D images. For 3D images, please refer to DeepMind's implementation + https://github.com/deepmind/surface-distance. + + The class- and batch sample-wise NSD values can be aggregated with the function `aggregate`. + + Args: + class_thresholds: List of class-specific thresholds. + The thresholds relate to the acceptable amount of deviation in the segmentation boundary in pixels. + Each threshold needs to be a finite, non-negative number. + include_background: Whether to skip NSD computation on the first channel of the predicted output. + Defaults to ``False``. + distance_metric: The metric used to compute surface distances. + One of [``"euclidean"``, ``"chessboard"``, ``"taxicab"``]. + Defaults to ``"euclidean"``. + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. + get_not_nans: whether to return the `not_nans` count. + Defaults to ``False``. + `not_nans` is the number of batch samples for which not all class-specific NSD values were nan values. + If set to ``True``, the function `aggregate` will return both the aggregated NSD and the `not_nans` count. + If set to ``False``, `aggregate` will only return the aggregated NSD. + """ + + def __init__( + self, + class_thresholds: List[float], + include_background: bool = False, + distance_metric: str = "euclidean", + reduction: Union[MetricReduction, str] = MetricReduction.MEAN, + get_not_nans: bool = False, + ) -> None: + super().__init__() + self.class_thresholds = class_thresholds + self.include_background = include_background + self.distance_metric = distance_metric + self.reduction = reduction + self.get_not_nans = get_not_nans + + def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignore + r""" + Args: + y_pred: Predicted segmentation, typically segmentation model output. + It must be a one-hot encoded, batch-first tensor [B,C,H,W]. + y: Reference segmentation. + It must be a one-hot encoded, batch-first tensor [B,C,H,W]. + + Returns: + Pytorch Tensor of shape [B,C], containing the NSD values :math:`\operatorname {NSD}_{b,c}` for each batch + index :math:`b` and class :math:`c`. + """ + return compute_surface_dice( + y_pred=y_pred, + y=y, + class_thresholds=self.class_thresholds, + include_background=self.include_background, + distance_metric=self.distance_metric, + ) + + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore + r""" + Aggregates the output of `_compute_tensor`. + + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + + Returns: + If `get_not_nans` is set to ``True``, this function returns the aggregated NSD and the `not_nans` count. + If `get_not_nans` is set to ``False``, this function returns only the aggregated NSD. + """ + data = self.get_buffer() + if not isinstance(data, torch.Tensor): + raise ValueError("the data to aggregate must be PyTorch Tensor.") + + # do metric reduction + f, not_nans = do_metric_reduction(data, reduction or self.reduction) + return (f, not_nans) if self.get_not_nans else f + + +def compute_surface_dice( + y_pred: torch.Tensor, + y: torch.Tensor, + class_thresholds: List[float], + include_background: bool = False, + distance_metric: str = "euclidean", +): + r""" + This function computes the (Normalized) Surface Dice (NSD) between the two tensors `y_pred` (referred to as + :math:`\hat{Y}`) and `y` (referred to as :math:`Y`). This metric determines which fraction of a segmentation + boundary is correctly predicted. A boundary element is considered correctly predicted if the closest distance to the + reference boundary is smaller than or equal to the specified threshold related to the acceptable amount of deviation in + pixels. The NSD is bounded between 0 and 1. + + This implementation supports multi-class tasks with an individual threshold :math:`\tau_c` for each class :math:`c`. + The class-specific NSD for batch index :math:`b`, :math:`\operatorname {NSD}_{b,c}`, is computed using the function: + + .. math:: + \operatorname {NSD}_{b,c} \left(Y_{b,c}, \hat{Y}_{b,c}\right) = \frac{\left|\mathcal{D}_{Y_{b,c}}^{'}\right| + + \left| \mathcal{D}_{\hat{Y}_{b,c}}^{'} \right|}{\left|\mathcal{D}_{Y_{b,c}}\right| + + \left|\mathcal{D}_{\hat{Y}_{b,c}}\right|} + :label: nsd + + with :math:`\mathcal{D}_{Y_{b,c}}` and :math:`\mathcal{D}_{\hat{Y}_{b,c}}` being two sets of nearest-neighbor + distances. :math:`\mathcal{D}_{Y_{b,c}}` is computed from the predicted segmentation boundary towards the reference segmentation + boundary and vice-versa for :math:`\mathcal{D}_{\hat{Y}_{b,c}}`. :math:`\mathcal{D}_{Y_{b,c}}^{'}` and + :math:`\mathcal{D}_{\hat{Y}_{b,c}}^{'}` refer to the subsets of distances that are smaller or equal to the + acceptable distance :math:`\tau_c`: + + .. math:: + \mathcal{D}_{Y_{b,c}}^{'} = \{ d \in \mathcal{D}_{Y_{b,c}} \, | \, d \leq \tau_c \}. + + + In the case of a class neither being present in the predicted segmentation, nor in the reference segmentation, a nan value + will be returned for this class. In the case of a class being present in only one of predicted segmentation or + reference segmentation, the class NSD will be 0. + + This implementation is based on https://arxiv.org/abs/2111.05408 and supports 2D images. + Be aware that the computation of boundaries is different from DeepMind's implementation + https://github.com/deepmind/surface-distance. In this implementation, the length of a segmentation boundary is + interpreted as the number of its edge pixels. In DeepMind's implementation, the length of a segmentation boundary + depends on the local neighborhood (cf. https://arxiv.org/abs/1809.04430). + + Args: + y_pred: Predicted segmentation, typically segmentation model output. + It must be a one-hot encoded, batch-first tensor [B,C,H,W]. + y: Reference segmentation. + It must be a one-hot encoded, batch-first tensor [B,C,H,W]. + class_thresholds: List of class-specific thresholds. + The thresholds relate to the acceptable amount of deviation in the segmentation boundary in pixels. + Each threshold needs to be a finite, non-negative number. + include_background: Whether to skip the surface dice computation on the first channel of + the predicted output. Defaults to ``False``. + distance_metric: The metric used to compute surface distances. + One of [``"euclidean"``, ``"chessboard"``, ``"taxicab"``]. + Defaults to ``"euclidean"``. + + Raises: + ValueError: If `y_pred` and/or `y` are not PyTorch tensors. + ValueError: If `y_pred` and/or `y` do not have four dimensions. + ValueError: If `y_pred` and/or `y` have different shapes. + ValueError: If `y_pred` and/or `y` are not one-hot encoded + ValueError: If the number of channels of `y_pred` and/or `y` is different from the number of class thresholds. + ValueError: If any class threshold is not finite. + ValueError: If any class threshold is negative. + + Returns: + Pytorch Tensor of shape [B,C], containing the NSD values :math:`\operatorname {NSD}_{b,c}` for each batch index + :math:`b` and class :math:`c`. + """ + + if not include_background: + y_pred, y = ignore_background(y_pred=y_pred, y=y) + + if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): + raise ValueError("y_pred and y must be PyTorch Tensor.") + + if y_pred.ndimension() != 4 or y.ndimension() != 4: + raise ValueError("y_pred and y should have four dimensions: [B,C,H,W].") + + if y_pred.shape != y.shape: + raise ValueError( + f"y_pred and y should have same shape, but instead, shapes are {y_pred.shape} (y_pred) and {y.shape} (y)." + ) + + if not torch.all(y_pred.byte() == y_pred) or not torch.all(y.byte() == y): + raise ValueError("y_pred and y should be binarized tensors (e.g. torch.int64).") + if torch.any(y_pred > 1) or torch.any(y > 1): + raise ValueError("y_pred and y should be one-hot encoded.") + + y = y.float() + y_pred = y_pred.float() + + batch_size, n_class = y_pred.shape[:2] + + if n_class != len(class_thresholds): + raise ValueError( + f"number of classes ({n_class}) does not match number of class thresholds ({len(class_thresholds)})." + ) + + if any(~np.isfinite(class_thresholds)): + raise ValueError("All class thresholds need to be finite.") + + if any(np.array(class_thresholds) < 0): + raise ValueError("All class thresholds need to be >= 0.") + + nsd = np.empty((batch_size, n_class)) + + for b, c in np.ndindex(batch_size, n_class): + (edges_pred, edges_gt) = get_mask_edges(y_pred[b, c], y[b, c], crop=False) + if not np.any(edges_gt): + warnings.warn(f"the ground truth of class {c} is all 0, this may result in nan/inf distance.") + if not np.any(edges_pred): + warnings.warn(f"the prediction of class {c} is all 0, this may result in nan/inf distance.") + + distances_pred_gt = get_surface_distance(edges_pred, edges_gt, distance_metric=distance_metric) + distances_gt_pred = get_surface_distance(edges_gt, edges_pred, distance_metric=distance_metric) + + boundary_complete = len(distances_pred_gt) + len(distances_gt_pred) + boundary_correct = np.sum(distances_pred_gt <= class_thresholds[c]) + np.sum( + distances_gt_pred <= class_thresholds[c] + ) + + if boundary_complete == 0: + # the class is neither present in the prediction, nor in the reference segmentation + nsd[b, c] = np.nan + else: + nsd[b, c] = boundary_correct / boundary_complete + + return convert_data_type(nsd, torch.Tensor)[0] diff --git a/monai/metrics/surface_distance.py b/monai/metrics/surface_distance.py index 2c84bb9e7c..4b5cdbc2f7 100644 --- a/monai/metrics/surface_distance.py +++ b/monai/metrics/surface_distance.py @@ -15,7 +15,13 @@ import numpy as np import torch -from monai.metrics.utils import do_metric_reduction, get_mask_edges, get_surface_distance, ignore_background +from monai.metrics.utils import ( + do_metric_reduction, + get_mask_edges, + get_surface_distance, + ignore_background, + is_binary_tensor, +) from monai.utils import MetricReduction, convert_data_type from .metric import CumulativeIterationMetric @@ -37,7 +43,7 @@ class SurfaceDistanceMetric(CumulativeIterationMetric): `seg_pred` and `seg_gt`. Defaults to ``False``. distance_metric: : [``"euclidean"``, ``"chessboard"``, ``"taxicab"``] the metric used to compute surface distance. Defaults to ``"euclidean"``. - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction. get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans). @@ -73,14 +79,9 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor ValueError: when `y` is not a binarized tensor. ValueError: when `y_pred` has less than three dimensions. """ - if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): - raise ValueError("y_pred and y must be PyTorch Tensor.") - if not torch.all(y_pred.byte() == y_pred): - warnings.warn("y_pred should be a binarized tensor.") - if not torch.all(y.byte() == y): - warnings.warn("y should be a binarized tensor.") - dims = y_pred.ndimension() - if dims < 3: + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + if y_pred.dim() < 3: raise ValueError("y_pred should have at least three dimensions.") # compute (BxC) for each channel for each batch return compute_average_surface_distance( @@ -91,17 +92,22 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignor distance_metric=self.distance_metric, ) - def aggregate(self): + def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore """ Execute reduction logic for the output of `compute_average_surface_distance`. + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values, + available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, + ``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction. + """ data = self.get_buffer() if not isinstance(data, torch.Tensor): raise ValueError("the data to aggregate must be PyTorch Tensor.") # do metric reduction - f, not_nans = do_metric_reduction(data, self.reduction) + f, not_nans = do_metric_reduction(data, reduction or self.reduction) return (f, not_nans) if self.get_not_nans else f diff --git a/monai/metrics/utils.py b/monai/metrics/utils.py index fc42100d6f..c17df7a54a 100644 --- a/monai/metrics/utils.py +++ b/monai/metrics/utils.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from typing import Tuple, Union import numpy as np @@ -16,13 +17,13 @@ from monai.transforms.croppad.array import SpatialCrop from monai.transforms.utils import generate_spatial_bounding_box -from monai.utils import MetricReduction, look_up_option, optional_import +from monai.utils import MetricReduction, convert_data_type, look_up_option, optional_import binary_erosion, _ = optional_import("scipy.ndimage.morphology", name="binary_erosion") distance_transform_edt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_edt") distance_transform_cdt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_cdt") -__all__ = ["ignore_background", "do_metric_reduction", "get_mask_edges", "get_surface_distance"] +__all__ = ["ignore_background", "do_metric_reduction", "get_mask_edges", "get_surface_distance", "is_binary_tensor"] def ignore_background(y_pred: Union[np.ndarray, torch.Tensor], y: Union[np.ndarray, torch.Tensor]): @@ -50,11 +51,10 @@ def do_metric_reduction(f: torch.Tensor, reduction: Union[MetricReduction, str] Args: f: a tensor that contains the calculated metric scores per batch and per class. The first two dims should be batch and class. - reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values, + reduction: define the mode to reduce metrics, will only apply reduction on `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", return the input f tensor and not_nans. - Define the mode to reduce computation result of 1 batch data. Defaults to ``"mean"``. Raises: ValueError: When ``reduction`` is not one of @@ -103,12 +103,7 @@ def do_metric_reduction(f: torch.Tensor, reduction: Union[MetricReduction, str] return f, not_nans -def get_mask_edges( - seg_pred: Union[np.ndarray, torch.Tensor], - seg_gt: Union[np.ndarray, torch.Tensor], - label_idx: int = 1, - crop: bool = True, -) -> Tuple[np.ndarray, np.ndarray]: +def get_mask_edges(seg_pred, seg_gt, label_idx: int = 1, crop: bool = True) -> Tuple[np.ndarray, np.ndarray]: """ Do binary erosion and use XOR for input to get the edges. This function is helpful to further calculate metrics such as Average Surface @@ -156,10 +151,12 @@ def get_mask_edges( if not np.any(seg_pred | seg_gt): return np.zeros_like(seg_pred), np.zeros_like(seg_gt) - seg_pred, seg_gt = np.expand_dims(seg_pred, 0), np.expand_dims(seg_gt, 0) + channel_dim = 0 + seg_pred, seg_gt = np.expand_dims(seg_pred, axis=channel_dim), np.expand_dims(seg_gt, axis=channel_dim) box_start, box_end = generate_spatial_bounding_box(np.asarray(seg_pred | seg_gt)) cropper = SpatialCrop(roi_start=box_start, roi_end=box_end) - seg_pred, seg_gt = np.squeeze(cropper(seg_pred)), np.squeeze(cropper(seg_gt)) + seg_pred = convert_data_type(np.squeeze(cropper(seg_pred), axis=channel_dim), np.ndarray)[0] + seg_gt = convert_data_type(np.squeeze(cropper(seg_gt), axis=channel_dim), np.ndarray)[0] # Do binary erosion and use XOR to get edges edges_pred = binary_erosion(seg_pred) ^ seg_pred @@ -201,3 +198,22 @@ def get_surface_distance(seg_pred: np.ndarray, seg_gt: np.ndarray, distance_metr raise ValueError(f"distance_metric {distance_metric} is not implemented.") return np.asarray(dis[seg_pred]) + + +def is_binary_tensor(input: torch.Tensor, name: str): + """Determines whether the input tensor is torch binary tensor or not. + + Args: + input (torch.Tensor): tensor to validate. + name (str): name of the tensor being checked. + + Raises: + ValueError: if `input` is not a PyTorch Tensor. + + Returns: + Union[str, None]: warning message, if the tensor is not binary. Othwerwise, 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.") diff --git a/monai/networks/__init__.py b/monai/networks/__init__.py index 76223dfaef..0543b11632 100644 --- a/monai/networks/__init__.py +++ b/monai/networks/__init__.py @@ -20,6 +20,8 @@ one_hot, pixelshuffle, predict_segmentation, + replace_modules, + replace_modules_temp, save_state, slice_channels, to_norm_affine, diff --git a/monai/networks/blocks/__init__.py b/monai/networks/blocks/__init__.py index 0fdc944760..27feffea10 100644 --- a/monai/networks/blocks/__init__.py +++ b/monai/networks/blocks/__init__.py @@ -12,15 +12,17 @@ from .acti_norm import ADN from .activation import MemoryEfficientSwish, Mish, Swish from .aspp import SimpleASPP +from .backbone_fpn_utils import BackboneWithFPN from .convolutions import Convolution, ResidualUnit from .crf import CRF from .dints_block import ActiConvNormBlock, FactorizedIncreaseBlock, FactorizedReduceBlock, P3DActiConvNormBlock from .downsample import MaxAvgPool from .dynunet_block import UnetBasicBlock, UnetOutBlock, UnetResBlock, UnetUpBlock, get_output_padding, get_padding from .fcn import FCN, GCN, MCFCN, Refine +from .feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool, LastLevelP6P7 from .localnet_block import LocalNetDownSampleBlock, LocalNetFeatureExtractorBlock, LocalNetUpSampleBlock from .mlp import MLPBlock -from .patchembedding import PatchEmbeddingBlock +from .patchembedding import PatchEmbed, PatchEmbeddingBlock from .regunet_block import RegistrationDownSampleBlock, RegistrationExtractionBlock, RegistrationResidualConvBlock from .segresnet_block import ResBlock from .selfattention import SABlock diff --git a/monai/networks/blocks/acti_norm.py b/monai/networks/blocks/acti_norm.py index 65b662ac32..6aeaa7d275 100644 --- a/monai/networks/blocks/acti_norm.py +++ b/monai/networks/blocks/acti_norm.py @@ -18,8 +18,8 @@ class ADN(nn.Sequential): """ - Constructs a sequential module of optional activation, dropout, and normalization layers - (with an arbitrary order):: + Constructs a sequential module of optional activation (A), dropout (D), and normalization (N) layers + with an arbitrary order:: -- (Norm) -- (Dropout) -- (Acti) -- diff --git a/monai/networks/blocks/backbone_fpn_utils.py b/monai/networks/blocks/backbone_fpn_utils.py new file mode 100644 index 0000000000..c663485583 --- /dev/null +++ b/monai/networks/blocks/backbone_fpn_utils.py @@ -0,0 +1,176 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE +# +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +This script is modified from from torchvision to support N-D images, +by overriding the definition of convolutional layers and pooling layers. + +https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py +""" + +from typing import Dict, List, Optional, Union + +from torch import Tensor, nn + +from monai.networks.nets import resnet +from monai.utils import optional_import + +from .feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool + +torchvision_models, _ = optional_import("torchvision.models") + +__all__ = ["BackboneWithFPN"] + + +class BackboneWithFPN(nn.Module): + """ + Adds an FPN on top of a model. + Internally, it uses torchvision.models._utils.IntermediateLayerGetter to + extract a submodel that returns the feature maps specified in return_layers. + The same limitations of IntermediateLayerGetter apply here. + + Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py + Except that this class uses spatial_dims + + Args: + backbone: backbone network + return_layers: a dict containing the names + of the modules for which the activations will be returned as + the key of the dict, and the value of the dict is the name + of the returned activation (which the user can specify). + in_channels_list: number of channels for each feature map + that is returned, in the order they are present in the OrderedDict + out_channels: number of channels in the FPN. + spatial_dims: 2D or 3D images + """ + + def __init__( + self, + backbone: nn.Module, + return_layers: Dict[str, str], + in_channels_list: List[int], + out_channels: int, + spatial_dims: Union[int, None] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, + ) -> None: + super().__init__() + + # if spatial_dims is not specified, try to find it from backbone. + if spatial_dims is None: + if hasattr(backbone, "spatial_dims") and isinstance(backbone.spatial_dims, int): + spatial_dims = backbone.spatial_dims + elif isinstance(backbone.conv1, nn.Conv2d): + spatial_dims = 2 + elif isinstance(backbone.conv1, nn.Conv3d): + spatial_dims = 3 + else: + raise ValueError("Could not find spatial_dims of backbone, please specify it.") + + if extra_blocks is None: + extra_blocks = LastLevelMaxPool(spatial_dims) + + self.body = torchvision_models._utils.IntermediateLayerGetter(backbone, return_layers=return_layers) + self.fpn = FeaturePyramidNetwork( + spatial_dims=spatial_dims, + in_channels_list=in_channels_list, + out_channels=out_channels, + extra_blocks=extra_blocks, + ) + self.out_channels = out_channels + + def forward(self, x: Tensor) -> Dict[str, Tensor]: + """ + Computes the resulted feature maps of the network. + + Args: + x: input images + + Returns: + feature maps after FPN layers. They are ordered from highest resolution first. + """ + x = self.body(x) # backbone + y: Dict[str, Tensor] = self.fpn(x) # FPN + return y + + +def _resnet_fpn_extractor( + backbone: resnet.ResNet, + spatial_dims: int, + trainable_layers: int = 5, + returned_layers: Optional[List[int]] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, +) -> BackboneWithFPN: + """ + Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py + Except that ``in_channels_stage2 = backbone.in_planes // 8`` instead of ``in_channels_stage2 = backbone.inplanes // 8``, + and it requires spatial_dims: 2D or 3D images. + """ + + # select layers that wont be frozen + if trainable_layers < 0 or trainable_layers > 5: + raise ValueError(f"Trainable layers should be in the range [0,5], got {trainable_layers}") + layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers] + if trainable_layers == 5: + layers_to_train.append("bn1") + for name, parameter in backbone.named_parameters(): + if all([not name.startswith(layer) for layer in layers_to_train]): + parameter.requires_grad_(False) + + if extra_blocks is None: + extra_blocks = LastLevelMaxPool(spatial_dims) + + if returned_layers is None: + returned_layers = [1, 2, 3, 4] + if min(returned_layers) <= 0 or max(returned_layers) >= 5: + raise ValueError(f"Each returned layer should be in the range [1,4]. Got {returned_layers}") + return_layers = {f"layer{k}": str(v) for v, k in enumerate(returned_layers)} + + in_channels_stage2 = backbone.in_planes // 8 + in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers] + out_channels = 256 + return BackboneWithFPN( + backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks, spatial_dims=spatial_dims + ) diff --git a/monai/networks/blocks/convolutions.py b/monai/networks/blocks/convolutions.py index 37530668a3..17d3c9d084 100644 --- a/monai/networks/blocks/convolutions.py +++ b/monai/networks/blocks/convolutions.py @@ -159,19 +159,22 @@ def __init__( self.add_module("conv", conv) - if not conv_only: - self.add_module( - "adn", - ADN( - ordering=adn_ordering, - in_channels=out_channels, - act=act, - norm=norm, - norm_dim=self.dimensions, - dropout=dropout, - dropout_dim=dropout_dim, - ), - ) + if conv_only: + return + if act is None and norm is None and dropout is None: + return + self.add_module( + "adn", + ADN( + ordering=adn_ordering, + in_channels=out_channels, + act=act, + norm=norm, + norm_dim=self.dimensions, + dropout=dropout, + dropout_dim=dropout_dim, + ), + ) class ResidualUnit(nn.Module): diff --git a/monai/networks/blocks/dints_block.py b/monai/networks/blocks/dints_block.py index f76e125fe0..b7365f50e3 100644 --- a/monai/networks/blocks/dints_block.py +++ b/monai/networks/blocks/dints_block.py @@ -31,7 +31,7 @@ def __init__( out_channel: int, spatial_dims: int = 3, act_name: Union[Tuple, str] = "RELU", - norm_name: Union[Tuple, str] = "INSTANCE", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), ): """ Args: @@ -82,7 +82,7 @@ def __init__( out_channel: int, spatial_dims: int = 3, act_name: Union[Tuple, str] = "RELU", - norm_name: Union[Tuple, str] = "INSTANCE", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), ): """ Args: @@ -150,7 +150,7 @@ def __init__( padding: int, mode: int = 0, act_name: Union[Tuple, str] = "RELU", - norm_name: Union[Tuple, str] = "INSTANCE", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), ): """ Args: @@ -235,7 +235,7 @@ def __init__( padding: int = 1, spatial_dims: int = 3, act_name: Union[Tuple, str] = "RELU", - norm_name: Union[Tuple, str] = "INSTANCE", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), ): """ Args: diff --git a/monai/networks/blocks/dynunet_block.py b/monai/networks/blocks/dynunet_block.py index 8b22cb16a9..894686b50a 100644 --- a/monai/networks/blocks/dynunet_block.py +++ b/monai/networks/blocks/dynunet_block.py @@ -57,22 +57,41 @@ def __init__( kernel_size=kernel_size, stride=stride, dropout=dropout, - conv_only=True, + act=None, + norm=None, + conv_only=False, ) self.conv2 = get_conv_layer( - spatial_dims, out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, conv_only=True - ) - self.conv3 = get_conv_layer( - spatial_dims, in_channels, out_channels, kernel_size=1, stride=stride, dropout=dropout, conv_only=True + spatial_dims, + out_channels, + out_channels, + kernel_size=kernel_size, + stride=1, + dropout=dropout, + act=None, + norm=None, + conv_only=False, ) self.lrelu = get_act_layer(name=act_name) self.norm1 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels) self.norm2 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels) - self.norm3 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels) self.downsample = in_channels != out_channels stride_np = np.atleast_1d(stride) if not np.all(stride_np == 1): self.downsample = True + if self.downsample: + self.conv3 = get_conv_layer( + spatial_dims, + in_channels, + out_channels, + kernel_size=1, + stride=stride, + dropout=dropout, + act=None, + norm=None, + conv_only=False, + ) + self.norm3 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels) def forward(self, inp): residual = inp @@ -81,8 +100,9 @@ def forward(self, inp): out = self.lrelu(out) out = self.conv2(out) out = self.norm2(out) - if self.downsample: + if hasattr(self, "conv3"): residual = self.conv3(residual) + if hasattr(self, "norm3"): residual = self.norm3(residual) out += residual out = self.lrelu(out) @@ -91,7 +111,7 @@ def forward(self, inp): class UnetBasicBlock(nn.Module): """ - A CNN module module that can be used for DynUNet, based on: + A CNN module that can be used for DynUNet, based on: `Automated Design of Deep Learning Methods for Biomedical Image Segmentation `_. `nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation `_. @@ -126,10 +146,20 @@ def __init__( kernel_size=kernel_size, stride=stride, dropout=dropout, - conv_only=True, + act=None, + norm=None, + conv_only=False, ) self.conv2 = get_conv_layer( - spatial_dims, out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, conv_only=True + spatial_dims, + out_channels, + out_channels, + kernel_size=kernel_size, + stride=1, + dropout=dropout, + act=None, + norm=None, + conv_only=False, ) self.lrelu = get_act_layer(name=act_name) self.norm1 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels) @@ -188,7 +218,9 @@ def __init__( stride=upsample_stride, dropout=dropout, bias=trans_bias, - conv_only=True, + act=None, + norm=None, + conv_only=False, is_transposed=True, ) self.conv_block = UnetBasicBlock( @@ -216,7 +248,16 @@ def __init__( ): super().__init__() self.conv = get_conv_layer( - spatial_dims, in_channels, out_channels, kernel_size=1, stride=1, dropout=dropout, bias=True, conv_only=True + spatial_dims, + in_channels, + out_channels, + kernel_size=1, + stride=1, + dropout=dropout, + bias=True, + act=None, + norm=None, + conv_only=False, ) def forward(self, inp): @@ -230,7 +271,7 @@ def get_conv_layer( kernel_size: Union[Sequence[int], int] = 3, stride: Union[Sequence[int], int] = 1, act: Optional[Union[Tuple, str]] = Act.PRELU, - norm: Union[Tuple, str] = Norm.INSTANCE, + norm: Optional[Union[Tuple, str]] = Norm.INSTANCE, dropout: Optional[Union[Tuple, str, float]] = None, bias: bool = False, conv_only: bool = True, diff --git a/monai/networks/blocks/feature_pyramid_network.py b/monai/networks/blocks/feature_pyramid_network.py new file mode 100644 index 0000000000..2373cfc099 --- /dev/null +++ b/monai/networks/blocks/feature_pyramid_network.py @@ -0,0 +1,263 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/release/0.12/torchvision/ops/feature_pyramid_network.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE +# +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +This script is modified from from torchvision to support N-D images, +by overriding the definition of convolutional layers and pooling layers. + +https://github.com/pytorch/vision/blob/release/0.12/torchvision/ops/feature_pyramid_network.py +""" + +from collections import OrderedDict +from typing import Callable, Dict, List, Optional, Tuple, Type, Union + +import torch.nn.functional as F +from torch import Tensor, nn + +from monai.networks.layers.factories import Conv, Pool + +__all__ = ["ExtraFPNBlock", "LastLevelMaxPool", "LastLevelP6P7", "FeaturePyramidNetwork"] + + +class ExtraFPNBlock(nn.Module): + """ + Base class for the extra block in the FPN. + + Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/ops/feature_pyramid_network.py + """ + + def forward(self, results: List[Tensor], x: List[Tensor], names: List[str]) -> Tuple[List[Tensor], List[str]]: + """ + Compute extended set of results of the FPN and their names. + + Args: + results: the result of the FPN + x: the original feature maps + names: the names for each one of the original feature maps + + Returns: + - the extended set of results of the FPN + - the extended set of names for the results + """ + pass + + +class LastLevelMaxPool(ExtraFPNBlock): + """ + Applies a max_pool2d or max_pool3d on top of the last feature map. Serves as an ``extra_blocks`` + in :class:`~monai.networks.blocks.feature_pyramid_network.FeaturePyramidNetwork` . + """ + + def __init__(self, spatial_dims: int): + super().__init__() + pool_type: Type[Union[nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d]] = Pool[Pool.MAX, spatial_dims] + self.maxpool = pool_type(kernel_size=1, stride=2, padding=0) + + def forward(self, results: List[Tensor], x: List[Tensor], names: List[str]) -> Tuple[List[Tensor], List[str]]: + names.append("pool") + results.append(self.maxpool(results[-1])) + return results, names + + +class LastLevelP6P7(ExtraFPNBlock): + """ + This module is used in RetinaNet to generate extra layers, P6 and P7. + Serves as an ``extra_blocks`` + in :class:`~monai.networks.blocks.feature_pyramid_network.FeaturePyramidNetwork` . + """ + + def __init__(self, spatial_dims: int, in_channels: int, out_channels: int): + super().__init__() + conv_type: Callable = Conv[Conv.CONV, spatial_dims] + self.p6 = conv_type(in_channels, out_channels, kernel_size=3, stride=2, padding=1) + self.p7 = conv_type(out_channels, out_channels, kernel_size=3, stride=2, padding=1) + for module in [self.p6, self.p7]: + nn.init.kaiming_uniform_(module.weight, a=1) + nn.init.constant_(module.bias, 0) + self.use_P5 = in_channels == out_channels + + def forward(self, results: List[Tensor], x: List[Tensor], names: List[str]) -> Tuple[List[Tensor], List[str]]: + p5, c5 = results[-1], x[-1] + x5 = p5 if self.use_P5 else c5 + p6 = self.p6(x5) + p7 = self.p7(F.relu(p6)) + results.extend([p6, p7]) + names.extend(["p6", "p7"]) + return results, names + + +class FeaturePyramidNetwork(nn.Module): + """ + Module that adds a FPN from on top of a set of feature maps. This is based on + `"Feature Pyramid Network for Object Detection" `_. + + The feature maps are currently supposed to be in increasing depth + order. + + The input to the model is expected to be an OrderedDict[Tensor], containing + the feature maps on top of which the FPN will be added. + + Args: + spatial_dims: 2D or 3D images + in_channels_list: number of channels for each feature map that + is passed to the module + out_channels: number of channels of the FPN representation + extra_blocks: if provided, extra operations will + be performed. It is expected to take the fpn features, the original + features and the names of the original features as input, and returns + a new list of feature maps and their corresponding names + + Examples:: + + >>> m = FeaturePyramidNetwork(2, [10, 20, 30], 5) + >>> # get some dummy data + >>> x = OrderedDict() + >>> x['feat0'] = torch.rand(1, 10, 64, 64) + >>> x['feat2'] = torch.rand(1, 20, 16, 16) + >>> x['feat3'] = torch.rand(1, 30, 8, 8) + >>> # compute the FPN on top of x + >>> output = m(x) + >>> print([(k, v.shape) for k, v in output.items()]) + >>> # returns + >>> [('feat0', torch.Size([1, 5, 64, 64])), + >>> ('feat2', torch.Size([1, 5, 16, 16])), + >>> ('feat3', torch.Size([1, 5, 8, 8]))] + + """ + + def __init__( + self, + spatial_dims: int, + in_channels_list: List[int], + out_channels: int, + extra_blocks: Optional[ExtraFPNBlock] = None, + ): + super().__init__() + + conv_type: Callable = Conv[Conv.CONV, spatial_dims] + + self.inner_blocks = nn.ModuleList() + self.layer_blocks = nn.ModuleList() + for in_channels in in_channels_list: + if in_channels == 0: + raise ValueError("in_channels=0 is currently not supported") + inner_block_module = conv_type(in_channels, out_channels, 1) + layer_block_module = conv_type(out_channels, out_channels, 3, padding=1) + self.inner_blocks.append(inner_block_module) + self.layer_blocks.append(layer_block_module) + + # initialize parameters now to avoid modifying the initialization of top_blocks + conv_type_: Type[nn.Module] = Conv[Conv.CONV, spatial_dims] + for m in self.modules(): + if isinstance(m, conv_type_): + nn.init.kaiming_uniform_(m.weight, a=1) # type: ignore + nn.init.constant_(m.bias, 0.0) # type: ignore + + if extra_blocks is not None: + if not isinstance(extra_blocks, ExtraFPNBlock): + raise AssertionError + self.extra_blocks = extra_blocks + + def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: + """ + This is equivalent to self.inner_blocks[idx](x), + but torchscript doesn't support this yet + """ + num_blocks = len(self.inner_blocks) + if idx < 0: + idx += num_blocks + out = x + for i, module in enumerate(self.inner_blocks): + if i == idx: + out = module(x) + return out + + def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: + """ + This is equivalent to self.layer_blocks[idx](x), + but torchscript doesn't support this yet + """ + num_blocks = len(self.layer_blocks) + if idx < 0: + idx += num_blocks + out = x + for i, module in enumerate(self.layer_blocks): + if i == idx: + out = module(x) + return out + + def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: + """ + Computes the FPN for a set of feature maps. + + Args: + x: feature maps for each feature level. + + Returns: + feature maps after FPN layers. They are ordered from highest resolution first. + """ + # unpack OrderedDict into two lists for easier handling + names = list(x.keys()) + x_values: List[Tensor] = list(x.values()) + + last_inner = self.get_result_from_inner_blocks(x_values[-1], -1) + results = [] + results.append(self.get_result_from_layer_blocks(last_inner, -1)) + + for idx in range(len(x_values) - 2, -1, -1): + inner_lateral = self.get_result_from_inner_blocks(x_values[idx], idx) + feat_shape = inner_lateral.shape[2:] + inner_top_down = F.interpolate(last_inner, size=feat_shape, mode="nearest") + last_inner = inner_lateral + inner_top_down + results.insert(0, self.get_result_from_layer_blocks(last_inner, idx)) + + if self.extra_blocks is not None: + results, names = self.extra_blocks(results, x_values, names) + + # make it back an OrderedDict + out = OrderedDict([(k, v) for k, v in zip(names, results)]) + + return out diff --git a/monai/networks/blocks/fft_utils_t.py b/monai/networks/blocks/fft_utils_t.py new file mode 100644 index 0000000000..0d6b99d7e1 --- /dev/null +++ b/monai/networks/blocks/fft_utils_t.py @@ -0,0 +1,232 @@ +# 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 typing import Optional, Sequence + +import torch +from torch import Tensor + +from monai.utils.type_conversion import convert_data_type + + +def roll_1d(x: Tensor, shift: int, shift_dim: int) -> Tensor: + """ + Similar to roll but for only one dim. + + Args: + x: input data (k-space or image) that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + shift: the amount of shift along each of shift_dims dimension + shift_dim: the dimension over which the shift is applied + + Returns: + 1d-shifted version of x + + Note: + This function is called when fftshift and ifftshift are not available in the running pytorch version + """ + shift = shift % x.size(shift_dim) + if shift == 0: + return x + + left = x.narrow(shift_dim, 0, x.size(shift_dim) - shift) + right = x.narrow(shift_dim, x.size(shift_dim) - shift, shift) + + return torch.cat((right, left), dim=shift_dim) + + +def roll(x: Tensor, shift: Sequence[int], shift_dims: Sequence[int]) -> Tensor: + """ + Similar to np.roll but applies to PyTorch Tensors + + Args: + x: input data (k-space or image) that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + shift: the amount of shift along each of shift_dims dimensions + shift_dims: dimensions over which the shift is applied + + Returns: + shifted version of x + + Note: + This function is called when fftshift and ifftshift are not available in the running pytorch version + """ + if len(shift) != len(shift_dims): + raise ValueError(f"len(shift) != len(shift_dims), got f{len(shift)} and f{len(shift_dims)}.") + for s, d in zip(shift, shift_dims): + x = roll_1d(x, s, d) + return x + + +def fftshift(x: Tensor, shift_dims: Optional[Sequence[int]] = None) -> Tensor: + """ + Similar to np.fft.fftshift but applies to PyTorch Tensors + + Args: + x: input data (k-space or image) that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + shift_dims: dimensions over which the shift is applied + + Returns: + fft-shifted version of x + + Note: + This function is called when fftshift is not available in the running pytorch version + """ + if shift_dims is None: + # for torch.jit.script based on the fastmri repository + shift_dims = [0] * (x.dim()) + for i in range(1, x.dim()): + shift_dims[i] = i + shift = [0] * len(shift_dims) + for i, dim_num in enumerate(shift_dims): + shift[i] = x.shape[dim_num] // 2 + return roll(x, shift, shift_dims) + + +def ifftshift(x: Tensor, shift_dims: Optional[Sequence[int]] = None) -> Tensor: + """ + Similar to np.fft.ifftshift but applies to PyTorch Tensors + + Args: + x: input data (k-space or image) that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + shift_dims: dimensions over which the shift is applied + + Returns: + ifft-shifted version of x + + Note: + This function is called when ifftshift is not available in the running pytorch version + """ + if shift_dims is None: + # for torch.jit.script based on the fastmri repository + shift_dims = [0] * (x.dim()) + for i in range(1, x.dim()): + shift_dims[i] = i + shift = [0] * len(shift_dims) + for i, dim_num in enumerate(shift_dims): + shift[i] = (x.shape[dim_num] + 1) // 2 + return roll(x, shift, shift_dims) + + +def ifftn_centered_t(ksp: Tensor, spatial_dims: int, is_complex: bool = True) -> Tensor: + """ + Pytorch-based ifft for spatial_dims-dim signals. "centered" means this function automatically takes care + of the required ifft and fft shifts. + This is equivalent to do fft in numpy based on numpy.fft.ifftn, numpy.fft.fftshift, and numpy.fft.ifftshift + + Args: + ksp: k-space data that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) + is_complex: if True, then the last dimension of the input ksp is expected to be 2 (representing real and imaginary channels) + + Returns: + "out" which is the output image (inverse fourier of ksp) + + Example: + + .. code-block:: python + + import torch + ksp = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts + # output1 and output2 will be identical + output1 = torch.fft.ifftn(torch.view_as_complex(torch.fft.ifftshift(ksp,dim=(-3,-2))), dim=(-2,-1), norm="ortho") + output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) + + output2 = ifftn_centered(ksp, spatial_dims=2, is_complex=True) + """ + # define spatial dims to perform ifftshift, fftshift, and ifft + shift = tuple(range(-spatial_dims, 0)) + if is_complex: + if ksp.shape[-1] != 2: + raise ValueError(f"ksp.shape[-1] is not 2 ({ksp.shape[-1]}).") + shift = tuple(range(-spatial_dims - 1, -1)) + dims = tuple(range(-spatial_dims, 0)) + + # apply ifft + if hasattr(torch.fft, "ifftshift"): # ifftshift was added in pytorch 1.8 + x = torch.fft.ifftshift(ksp, dim=shift) + else: + x = ifftshift(ksp, shift) + + if is_complex: + x = torch.view_as_real(torch.fft.ifftn(torch.view_as_complex(x), dim=dims, norm="ortho")) + else: + x = torch.view_as_real(torch.fft.ifftn(x, dim=dims, norm="ortho")) + + if hasattr(torch.fft, "fftshift"): + out = convert_data_type(torch.fft.fftshift(x, dim=shift), torch.Tensor)[0] + else: + out = convert_data_type(fftshift(x, shift), torch.Tensor)[0] + + return out + + +def fftn_centered_t(im: Tensor, spatial_dims: int, is_complex: bool = True) -> Tensor: + """ + Pytorch-based fft for spatial_dims-dim signals. "centered" means this function automatically takes care + of the required ifft and fft shifts. + This is equivalent to do ifft in numpy based on numpy.fft.fftn, numpy.fft.fftshift, and numpy.fft.ifftshift + + Args: + im: image that can be + 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or + 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. + spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) + is_complex: if True, then the last dimension of the input im is expected to be 2 (representing real and imaginary channels) + + Returns: + "out" which is the output kspace (fourier of im) + + Example: + + .. code-block:: python + + import torch + im = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts + # output1 and output2 will be identical + output1 = torch.fft.fftn(torch.view_as_complex(torch.fft.ifftshift(im,dim=(-3,-2))), dim=(-2,-1), norm="ortho") + output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) + + output2 = fftn_centered(im, spatial_dims=2, is_complex=True) + """ + # define spatial dims to perform ifftshift, fftshift, and fft + shift = tuple(range(-spatial_dims, 0)) + if is_complex: + if im.shape[-1] != 2: + raise ValueError(f"img.shape[-1] is not 2 ({im.shape[-1]}).") + shift = tuple(range(-spatial_dims - 1, -1)) + dims = tuple(range(-spatial_dims, 0)) + + # apply fft + if hasattr(torch.fft, "ifftshift"): # ifftshift was added in pytorch 1.8 + x = torch.fft.ifftshift(im, dim=shift) + else: + x = ifftshift(im, shift) + + if is_complex: + x = torch.view_as_real(torch.fft.fftn(torch.view_as_complex(x), dim=dims, norm="ortho")) + else: + x = torch.view_as_real(torch.fft.fftn(x, dim=dims, norm="ortho")) + + if hasattr(torch.fft, "fftshift"): + out = convert_data_type(torch.fft.fftshift(x, dim=shift), torch.Tensor)[0] + else: + out = convert_data_type(fftshift(x, shift), torch.Tensor)[0] + + return out diff --git a/monai/networks/blocks/mlp.py b/monai/networks/blocks/mlp.py index a1728365cf..0feeb044f3 100644 --- a/monai/networks/blocks/mlp.py +++ b/monai/networks/blocks/mlp.py @@ -9,8 +9,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Tuple, Union + import torch.nn as nn +from monai.networks.layers import get_act_layer +from monai.utils import look_up_option + +SUPPORTED_DROPOUT_MODE = {"vit", "swin"} + class MLPBlock(nn.Module): """ @@ -18,12 +25,26 @@ class MLPBlock(nn.Module): An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " """ - def __init__(self, hidden_size: int, mlp_dim: int, dropout_rate: float = 0.0) -> None: + def __init__( + self, + hidden_size: int, + mlp_dim: int, + dropout_rate: float = 0.0, + act: Union[Tuple, str] = "GELU", + dropout_mode="vit", + ) -> None: """ Args: hidden_size: dimension of hidden layer. - mlp_dim: dimension of feedforward layer. + mlp_dim: dimension of feedforward layer. If 0, `hidden_size` will be used. dropout_rate: faction of the input units to drop. + act: activation type and arguments. Defaults to GELU. + dropout_mode: dropout mode, can be "vit" or "swin". + "vit" mode uses two dropout instances as implemented in + https://github.com/google-research/vision_transformer/blob/main/vit_jax/models.py#L87 + "swin" corresponds to one instance as implemented in + https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_mlp.py#L23 + """ @@ -31,12 +52,18 @@ def __init__(self, hidden_size: int, mlp_dim: int, dropout_rate: float = 0.0) -> if not (0 <= dropout_rate <= 1): raise ValueError("dropout_rate should be between 0 and 1.") - + mlp_dim = mlp_dim or hidden_size self.linear1 = nn.Linear(hidden_size, mlp_dim) self.linear2 = nn.Linear(mlp_dim, hidden_size) - self.fn = nn.GELU() + self.fn = get_act_layer(act) self.drop1 = nn.Dropout(dropout_rate) - self.drop2 = nn.Dropout(dropout_rate) + dropout_opt = look_up_option(dropout_mode, SUPPORTED_DROPOUT_MODE) + if dropout_opt == "vit": + self.drop2 = nn.Dropout(dropout_rate) + elif dropout_opt == "swin": + self.drop2 = self.drop1 + else: + raise ValueError(f"dropout_mode should be one of {SUPPORTED_DROPOUT_MODE}") def forward(self, x): x = self.fn(self.linear1(x)) diff --git a/monai/networks/blocks/patchembedding.py b/monai/networks/blocks/patchembedding.py index 4c7263c6d5..7dc84a9837 100644 --- a/monai/networks/blocks/patchembedding.py +++ b/monai/networks/blocks/patchembedding.py @@ -9,15 +9,15 @@ # See the License for the specific language governing permissions and # limitations under the License. - -import math -from typing import Sequence, Union +from typing import Sequence, Type, Union import numpy as np import torch import torch.nn as nn +import torch.nn.functional as F +from torch.nn import LayerNorm -from monai.networks.layers import Conv +from monai.networks.layers import Conv, trunc_normal_ from monai.utils import ensure_tuple_rep, optional_import from monai.utils.module import look_up_option @@ -98,34 +98,18 @@ def __init__( ) self.position_embeddings = nn.Parameter(torch.zeros(1, self.n_patches, hidden_size)) self.dropout = nn.Dropout(dropout_rate) - self.trunc_normal_(self.position_embeddings, mean=0.0, std=0.02, a=-2.0, b=2.0) + trunc_normal_(self.position_embeddings, mean=0.0, std=0.02, a=-2.0, b=2.0) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - self.trunc_normal_(m.weight, mean=0.0, std=0.02, a=-2.0, b=2.0) + trunc_normal_(m.weight, mean=0.0, std=0.02, a=-2.0, b=2.0) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) - def trunc_normal_(self, tensor, mean, std, a, b): - # From PyTorch official master until it's in a few official releases - RW - # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf - def norm_cdf(x): - return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 - - with torch.no_grad(): - l = norm_cdf((a - mean) / std) - u = norm_cdf((b - mean) / std) - tensor.uniform_(2 * l - 1, 2 * u - 1) - tensor.erfinv_() - tensor.mul_(std * math.sqrt(2.0)) - tensor.add_(mean) - tensor.clamp_(min=a, max=b) - return tensor - def forward(self, x): x = self.patch_embeddings(x) if self.pos_embed == "conv": @@ -133,3 +117,84 @@ def forward(self, x): embeddings = x + self.position_embeddings embeddings = self.dropout(embeddings) return embeddings + + +class PatchEmbed(nn.Module): + """ + Patch embedding block based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + + Unlike ViT patch embedding block: (1) input is padded to satisfy window size requirements (2) normalized if + specified (3) position embedding is not used. + + Example:: + + >>> from monai.networks.blocks import PatchEmbed + >>> PatchEmbed(patch_size=2, in_chans=1, embed_dim=48, norm_layer=nn.LayerNorm, spatial_dims=3) + """ + + def __init__( + self, + patch_size: Union[Sequence[int], int] = 2, + in_chans: int = 1, + embed_dim: int = 48, + norm_layer: Type[LayerNorm] = nn.LayerNorm, + spatial_dims: int = 3, + ) -> None: + """ + Args: + patch_size: dimension of patch size. + in_chans: dimension of input channels. + embed_dim: number of linear projection output channels. + norm_layer: normalization layer. + spatial_dims: spatial dimension. + """ + + super().__init__() + + if not (spatial_dims == 2 or spatial_dims == 3): + raise ValueError("spatial dimension should be 2 or 3.") + + patch_size = ensure_tuple_rep(patch_size, spatial_dims) + self.patch_size = patch_size + self.embed_dim = embed_dim + self.proj = Conv[Conv.CONV, spatial_dims]( + in_channels=in_chans, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size + ) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + x_shape = x.size() + if len(x_shape) == 5: + _, _, d, h, w = x_shape + if w % self.patch_size[2] != 0: + x = F.pad(x, (0, self.patch_size[2] - w % self.patch_size[2])) + if h % self.patch_size[1] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[1] - h % self.patch_size[1])) + if d % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - d % self.patch_size[0])) + + elif len(x_shape) == 4: + _, _, h, w = x.size() + if w % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - w % self.patch_size[1])) + if h % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - h % self.patch_size[0])) + + x = self.proj(x) + if self.norm is not None: + x_shape = x.size() + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + if len(x_shape) == 5: + d, wh, ww = x_shape[2], x_shape[3], x_shape[4] + x = x.transpose(1, 2).view(-1, self.embed_dim, d, wh, ww) + elif len(x_shape) == 4: + wh, ww = x_shape[2], x_shape[3] + x = x.transpose(1, 2).view(-1, self.embed_dim, wh, ww) + return x diff --git a/monai/networks/blocks/upsample.py b/monai/networks/blocks/upsample.py index fa3929df20..364db0e236 100644 --- a/monai/networks/blocks/upsample.py +++ b/monai/networks/blocks/upsample.py @@ -46,7 +46,7 @@ def __init__( size: Optional[Union[Tuple[int], int]] = None, mode: Union[UpsampleMode, str] = UpsampleMode.DECONV, pre_conv: Optional[Union[nn.Module, str]] = "default", - interp_mode: Union[InterpolateMode, str] = InterpolateMode.LINEAR, + interp_mode: str = InterpolateMode.LINEAR, align_corners: Optional[bool] = True, bias: bool = True, apply_pad_pool: bool = True, diff --git a/monai/networks/layers/__init__.py b/monai/networks/layers/__init__.py index 5115c00af3..f122dccee6 100644 --- a/monai/networks/layers/__init__.py +++ b/monai/networks/layers/__init__.py @@ -10,6 +10,7 @@ # limitations under the License. from .convutils import calculate_out_shape, gaussian_1d, polyval, same_padding, stride_minus_kernel_padding +from .drop_path import DropPath from .factories import Act, Conv, Dropout, LayerFactory, Norm, Pad, Pool, split_args from .filtering import BilateralFilter, PHLFilter from .gmm import GaussianMixtureModel @@ -27,3 +28,4 @@ ) from .spatial_transforms import AffineTransform, grid_count, grid_grad, grid_pull, grid_push from .utils import get_act_layer, get_dropout_layer, get_norm_layer, get_pool_layer +from .weight_init import _no_grad_trunc_normal_, trunc_normal_ diff --git a/monai/networks/layers/drop_path.py b/monai/networks/layers/drop_path.py new file mode 100644 index 0000000000..7bb209ed25 --- /dev/null +++ b/monai/networks/layers/drop_path.py @@ -0,0 +1,45 @@ +# 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. + +import torch.nn as nn + + +class DropPath(nn.Module): + """Stochastic drop paths per sample for residual blocks. + Based on: + https://github.com/rwightman/pytorch-image-models + """ + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True) -> None: + """ + Args: + drop_prob: drop path probability. + scale_by_keep: scaling by non-dropped probability. + """ + super().__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + if not (0 <= drop_prob <= 1): + raise ValueError("Drop path prob should be between 0 and 1.") + + def drop_path(self, x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + def forward(self, x): + return self.drop_path(x, self.drop_prob, self.training, self.scale_by_keep) diff --git a/monai/networks/layers/factories.py b/monai/networks/layers/factories.py index 6379f49449..89fe1912a5 100644 --- a/monai/networks/layers/factories.py +++ b/monai/networks/layers/factories.py @@ -60,11 +60,16 @@ def use_factory(fact_args): layer = use_factory( (fact.TEST, kwargs) ) """ +import warnings from typing import Any, Callable, Dict, Tuple, Type, Union +import torch import torch.nn as nn -from monai.utils import look_up_option +from monai.utils import look_up_option, optional_import + +InstanceNorm3dNVFuser, has_nvfuser = optional_import("apex.normalization", name="InstanceNorm3dNVFuser") + __all__ = ["LayerFactory", "Dropout", "Norm", "Act", "Conv", "Pool", "Pad", "split_args"] @@ -242,6 +247,43 @@ def sync_batch_factory(_dim) -> Type[nn.SyncBatchNorm]: return nn.SyncBatchNorm +@Norm.factory_function("instance_nvfuser") +def instance_nvfuser_factory(dim): + """ + `InstanceNorm3dNVFuser` is a faster verison of InstanceNorm layer and implemented in `apex`. + It only supports 3d tensors as the input. It also requires to use with CUDA and non-Windows OS. + In this function, if the required library `apex.normalization.InstanceNorm3dNVFuser` does not exist, + `nn.InstanceNorm3d` will be returned instead. + This layer is based on a customized autograd function, which is not supported in TorchScript currently. + Please switch to use `nn.InstanceNorm3d` if TorchScript is necessary. + + Please check the following link for more details about how to install `apex`: + https://github.com/NVIDIA/apex#installation + + """ + types = (nn.InstanceNorm1d, nn.InstanceNorm2d) + if dim != 3: + warnings.warn(f"`InstanceNorm3dNVFuser` only supports 3d cases, use {types[dim - 1]} instead.") + return types[dim - 1] + # test InstanceNorm3dNVFuser installation with a basic example + has_nvfuser_flag = has_nvfuser + if not torch.cuda.is_available(): + return nn.InstanceNorm3d + try: + layer = InstanceNorm3dNVFuser(num_features=1, affine=True).to("cuda:0") + inp = torch.randn([1, 1, 1, 1, 1]).to("cuda:0") + out = layer(inp) + del inp, out, layer + except Exception: + has_nvfuser_flag = False + if not has_nvfuser_flag: + warnings.warn( + "`apex.normalization.InstanceNorm3dNVFuser` is not installed properly, use nn.InstanceNorm3d instead." + ) + return nn.InstanceNorm3d + return InstanceNorm3dNVFuser + + Act.add_factory_callable("elu", lambda: nn.modules.ELU) Act.add_factory_callable("relu", lambda: nn.modules.ReLU) Act.add_factory_callable("leakyrelu", lambda: nn.modules.LeakyReLU) diff --git a/monai/networks/layers/simplelayers.py b/monai/networks/layers/simplelayers.py index 7a0a45cb64..3de4e75766 100644 --- a/monai/networks/layers/simplelayers.py +++ b/monai/networks/layers/simplelayers.py @@ -20,19 +20,11 @@ from monai.networks.layers.convutils import gaussian_1d from monai.networks.layers.factories import Conv -from monai.utils import ( - ChannelMatching, - InvalidPyTorchVersionError, - SkipMode, - look_up_option, - optional_import, - pytorch_after, -) +from monai.utils import ChannelMatching, SkipMode, look_up_option, optional_import, pytorch_after from monai.utils.misc import issequenceiterable _C, _ = optional_import("monai._C") -if pytorch_after(1, 7): - fft, _ = optional_import("torch.fft") +fft, _ = optional_import("torch.fft") __all__ = [ "ChannelPad", @@ -377,7 +369,6 @@ def _make_coeffs(window_length, order): class HilbertTransform(nn.Module): """ Determine the analytical signal of a Tensor along a particular axis. - Requires PyTorch 1.7.0+ and the PyTorch FFT module (which is not included in NVIDIA PyTorch Release 20.10). Args: axis: Axis along which to apply Hilbert transform. Default 2 (first spatial dimension). @@ -386,9 +377,6 @@ class HilbertTransform(nn.Module): def __init__(self, axis: int = 2, n: Union[int, None] = None) -> None: - if not pytorch_after(1, 7): - raise InvalidPyTorchVersionError("1.7.0", self.__class__.__name__) - super().__init__() self.axis = axis self.n = n diff --git a/monai/networks/layers/spatial_transforms.py b/monai/networks/layers/spatial_transforms.py index c1bb951c4d..7a41d79291 100644 --- a/monai/networks/layers/spatial_transforms.py +++ b/monai/networks/layers/spatial_transforms.py @@ -14,6 +14,7 @@ import torch import torch.nn as nn +import monai from monai.networks import to_norm_affine from monai.utils import GridSampleMode, GridSamplePadMode, ensure_tuple, look_up_option, optional_import @@ -116,6 +117,8 @@ def grid_pull( ] out: torch.Tensor out = _GridPull.apply(input, grid, interpolation, bound, extrapolate) + if isinstance(input, monai.data.MetaTensor): + out = monai.data.MetaTensor(out, meta=input.meta, applied_operations=input.applied_operations) return out @@ -217,7 +220,10 @@ def grid_push( if shape is None: shape = tuple(input.shape[2:]) - return _GridPush.apply(input, grid, shape, interpolation, bound, extrapolate) + out: torch.Tensor = _GridPush.apply(input, grid, shape, interpolation, bound, extrapolate) + if isinstance(input, monai.data.MetaTensor): + out = monai.data.MetaTensor(out, meta=input.meta, applied_operations=input.applied_operations) + return out class _GridCount(torch.autograd.Function): @@ -313,7 +319,10 @@ def grid_count(grid: torch.Tensor, shape=None, interpolation="linear", bound="ze if shape is None: shape = tuple(grid.shape[2:]) - return _GridCount.apply(grid, shape, interpolation, bound, extrapolate) + out: torch.Tensor = _GridCount.apply(grid, shape, interpolation, bound, extrapolate) + if isinstance(input, monai.data.MetaTensor): + out = monai.data.MetaTensor(out, meta=input.meta, applied_operations=input.applied_operations) + return out class _GridGrad(torch.autograd.Function): @@ -408,7 +417,10 @@ def grid_grad(input: torch.Tensor, grid: torch.Tensor, interpolation="linear", b for i in ensure_tuple(interpolation) ] - return _GridGrad.apply(input, grid, interpolation, bound, extrapolate) + out: torch.Tensor = _GridGrad.apply(input, grid, interpolation, bound, extrapolate) + if isinstance(input, monai.data.MetaTensor): + out = monai.data.MetaTensor(out, meta=input.meta, applied_operations=input.applied_operations) + return out class AffineTransform(nn.Module): @@ -416,10 +428,11 @@ def __init__( self, spatial_size: Optional[Union[Sequence[int], int]] = None, normalized: bool = False, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.ZEROS, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.ZEROS, align_corners: bool = False, reverse_indexing: bool = True, + zero_centered: Optional[bool] = None, ) -> None: """ Apply affine transformations with a batch of affine matrices. @@ -454,14 +467,23 @@ def __init__( reverse_indexing: whether to reverse the spatial indexing of image and coordinates. set to `False` if `theta` follows pytorch's default "D, H, W" convention. set to `True` if `theta` follows `scipy.ndimage` default "i, j, k" convention. + zero_centered: whether the affine is applied to coordinates in a zero-centered value range. + With `zero_centered=True`, for example, the center of rotation will be the + spatial center of the input; with `zero_centered=False`, the center of rotation will be the + origin of the input. This option is only available when `normalized=False`, + where the default behaviour is `False` if unspecified. + See also: :py:func:`monai.networks.utils.normalize_transform`. """ super().__init__() self.spatial_size = ensure_tuple(spatial_size) if spatial_size is not None else None self.normalized = normalized - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) self.align_corners = align_corners self.reverse_indexing = reverse_indexing + if zero_centered is not None and self.normalized: + raise ValueError("`normalized=True` is not compatible with the `zero_centered` option.") + self.zero_centered = zero_centered if zero_centered is not None else False def forward( self, src: torch.Tensor, theta: torch.Tensor, spatial_size: Optional[Union[Sequence[int], int]] = None @@ -525,7 +547,11 @@ def forward( # reverse and normalize theta if needed if not self.normalized: theta = to_norm_affine( - affine=theta, src_size=src_size[2:], dst_size=dst_size[2:], align_corners=self.align_corners + affine=theta, + src_size=src_size[2:], + dst_size=dst_size[2:], + align_corners=self.align_corners, + zero_centered=self.zero_centered, ) if self.reverse_indexing: rev_idx = torch.as_tensor(range(sr - 1, -1, -1), device=src.device) @@ -543,8 +569,8 @@ def forward( dst = nn.functional.grid_sample( input=src.contiguous(), grid=grid, - mode=self.mode.value, - padding_mode=self.padding_mode.value, + mode=self.mode, + padding_mode=self.padding_mode, align_corners=self.align_corners, ) return dst diff --git a/monai/networks/layers/weight_init.py b/monai/networks/layers/weight_init.py new file mode 100644 index 0000000000..9b81ef17f8 --- /dev/null +++ b/monai/networks/layers/weight_init.py @@ -0,0 +1,64 @@ +# 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. + +import math + +import torch + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + """Tensor initialization with truncated normal distribution. + Based on: + https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + https://github.com/rwightman/pytorch-image-models + + Args: + tensor: an n-dimensional `torch.Tensor`. + mean: the mean of the normal distribution. + std: the standard deviation of the normal distribution. + a: the minimum cutoff value. + b: the maximum cutoff value. + """ + + def norm_cdf(x): + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + with torch.no_grad(): + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + tensor.uniform_(2 * l - 1, 2 * u - 1) + tensor.erfinv_() + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + tensor.clamp_(min=a, max=b) + return tensor + + +def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0): + """Tensor initialization with truncated normal distribution. + Based on: + https://github.com/rwightman/pytorch-image-models + + Args: + tensor: an n-dimensional `torch.Tensor` + mean: the mean of the normal distribution + std: the standard deviation of the normal distribution + a: the minimum cutoff value + b: the maximum cutoff value + """ + + if not std > 0: + raise ValueError("the standard deviation should be greater than zero.") + + if a >= b: + raise ValueError("minimum cutoff value (a) should be smaller than maximum cutoff value (b).") + + return _no_grad_trunc_normal_(tensor, mean, std, a, b) diff --git a/monai/networks/nets/__init__.py b/monai/networks/nets/__init__.py index 16686fa25c..6bf1868122 100644 --- a/monai/networks/nets/__init__.py +++ b/monai/networks/nets/__init__.py @@ -43,6 +43,7 @@ from .fullyconnectednet import FullyConnectedNet, VarFullyConnectedNet from .generator import Generator from .highresnet import HighResBlock, HighResNet +from .hovernet import Hovernet, HoVernet, HoVerNet, HoverNet from .milmodel import MILModel from .netadapter import NetAdapter from .regressor import Regressor @@ -80,7 +81,8 @@ seresnext50, seresnext101, ) -from .torchvision_fc import TorchVisionFCModel, TorchVisionFullyConvModel +from .swin_unetr import SwinUNETR +from .torchvision_fc import TorchVisionFCModel from .transchex import BertAttention, BertMixedLayer, BertOutput, BertPreTrainedModel, MultiModal, Pooler, Transchex from .unet import UNet, Unet from .unetr import UNETR diff --git a/monai/networks/nets/dints.py b/monai/networks/nets/dints.py index 978695c5d0..b7f3921a47 100644 --- a/monai/networks/nets/dints.py +++ b/monai/networks/nets/dints.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + import warnings from typing import List, Optional, Tuple, Union @@ -96,22 +97,47 @@ class _ActiConvNormBlockWithRAMCost(ActiConvNormBlock): ram_cost = total_ram/output_size = 2 * in_channel/out_channel + 1 """ - def __init__(self, in_channel: int, out_channel: int, kernel_size: int, padding: int, spatial_dims: int = 3): - super().__init__(in_channel, out_channel, kernel_size, padding, spatial_dims) + def __init__( + self, + in_channel: int, + out_channel: int, + kernel_size: int, + padding: int, + spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), + ): + super().__init__(in_channel, out_channel, kernel_size, padding, spatial_dims, act_name, norm_name) self.ram_cost = 1 + in_channel / out_channel * 2 class _P3DActiConvNormBlockWithRAMCost(P3DActiConvNormBlock): - def __init__(self, in_channel: int, out_channel: int, kernel_size: int, padding: int, p3dmode: int = 0): - super().__init__(in_channel, out_channel, kernel_size, padding, p3dmode) + def __init__( + self, + in_channel: int, + out_channel: int, + kernel_size: int, + padding: int, + p3dmode: int = 0, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), + ): + super().__init__(in_channel, out_channel, kernel_size, padding, p3dmode, act_name, norm_name) # 1 in_channel (activation) + 1 in_channel (convolution) + # 1 out_channel (convolution) + 1 out_channel (normalization) self.ram_cost = 2 + 2 * in_channel / out_channel class _FactorizedIncreaseBlockWithRAMCost(FactorizedIncreaseBlock): - def __init__(self, in_channel: int, out_channel: int, spatial_dims: int = 3): - super().__init__(in_channel, out_channel, spatial_dims) + def __init__( + self, + in_channel: int, + out_channel: int, + spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), + ): + super().__init__(in_channel, out_channel, spatial_dims, act_name, norm_name) # s0 is upsampled 2x from s1, representing feature sizes at two resolutions. # 2 * in_channel * s0 (upsample + activation) + 2 * out_channel * s0 (conv + normalization) # s0 = output_size/out_channel @@ -119,8 +145,15 @@ def __init__(self, in_channel: int, out_channel: int, spatial_dims: int = 3): class _FactorizedReduceBlockWithRAMCost(FactorizedReduceBlock): - def __init__(self, in_channel: int, out_channel: int, spatial_dims: int = 3): - super().__init__(in_channel, out_channel, spatial_dims) + def __init__( + self, + in_channel: int, + out_channel: int, + spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), + ): + super().__init__(in_channel, out_channel, spatial_dims, act_name, norm_name) # s0 is upsampled 2x from s1, representing feature sizes at two resolutions. # in_channel * s0 (activation) + 3 * out_channel * s1 (convolution, concatenation, normalization) # s0 = s1 * 2^(spatial_dims) = output_size / out_channel * 2^(spatial_dims) @@ -182,14 +215,6 @@ class Cell(CellInterface): # \ # - Downsample - # Define connection operation set, parameterized by the number of channels - ConnOPS = { - "up": _FactorizedIncreaseBlockWithRAMCost, - "down": _FactorizedReduceBlockWithRAMCost, - "identity": _IdentityWithRAMCost, - "align_channels": _ActiConvNormBlockWithRAMCost, - } - # Define 2D operation set, parameterized by the number of channels OPS2D = { "skip_connect": lambda _c: _IdentityWithRAMCost(), @@ -205,18 +230,69 @@ class Cell(CellInterface): "conv_1x3x3": lambda c: _P3DActiConvNormBlockWithRAMCost(c, c, 3, padding=1, p3dmode=2), } - def __init__(self, c_prev: int, c: int, rate: int, arch_code_c=None, spatial_dims: int = 3): + # Define connection operation set, parameterized by the number of channels + ConnOPS = { + "up": _FactorizedIncreaseBlockWithRAMCost, + "down": _FactorizedReduceBlockWithRAMCost, + "identity": _IdentityWithRAMCost, + "align_channels": _ActiConvNormBlockWithRAMCost, + } + + def __init__( + self, + c_prev: int, + c: int, + rate: int, + arch_code_c=None, + spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), + ): super().__init__() self._spatial_dims = spatial_dims + self._act_name = act_name + self._norm_name = norm_name + if rate == -1: # downsample - self.preprocess = self.ConnOPS["down"](c_prev, c, spatial_dims=self._spatial_dims) + self.preprocess = self.ConnOPS["down"]( + c_prev, c, spatial_dims=self._spatial_dims, act_name=self._act_name, norm_name=self._norm_name + ) elif rate == 1: # upsample - self.preprocess = self.ConnOPS["up"](c_prev, c, spatial_dims=self._spatial_dims) + self.preprocess = self.ConnOPS["up"]( + c_prev, c, spatial_dims=self._spatial_dims, act_name=self._act_name, norm_name=self._norm_name + ) else: if c_prev == c: self.preprocess = self.ConnOPS["identity"]() else: - self.preprocess = self.ConnOPS["align_channels"](c_prev, c, 1, 0, spatial_dims=self._spatial_dims) + self.preprocess = self.ConnOPS["align_channels"]( + c_prev, c, 1, 0, spatial_dims=self._spatial_dims, act_name=self._act_name, norm_name=self._norm_name + ) + + # Define 2D operation set, parameterized by the number of channels + self.OPS2D = { + "skip_connect": lambda _c: _IdentityWithRAMCost(), + "conv_3x3": lambda c: _ActiConvNormBlockWithRAMCost( + c, c, 3, padding=1, spatial_dims=2, act_name=self._act_name, norm_name=self._norm_name + ), + } + + # Define 3D operation set, parameterized by the number of channels + self.OPS3D = { + "skip_connect": lambda _c: _IdentityWithRAMCost(), + "conv_3x3x3": lambda c: _ActiConvNormBlockWithRAMCost( + c, c, 3, padding=1, spatial_dims=3, act_name=self._act_name, norm_name=self._norm_name + ), + "conv_3x3x1": lambda c: _P3DActiConvNormBlockWithRAMCost( + c, c, 3, padding=1, p3dmode=0, act_name=self._act_name, norm_name=self._norm_name + ), + "conv_3x1x3": lambda c: _P3DActiConvNormBlockWithRAMCost( + c, c, 3, padding=1, p3dmode=1, act_name=self._act_name, norm_name=self._norm_name + ), + "conv_1x3x3": lambda c: _P3DActiConvNormBlockWithRAMCost( + c, c, 3, padding=1, p3dmode=2, act_name=self._act_name, norm_name=self._norm_name + ), + } self.OPS = {} if self._spatial_dims == 2: @@ -283,7 +359,7 @@ def __init__( in_channels: int, num_classes: int, act_name: Union[Tuple, str] = "RELU", - norm_name: Union[Tuple, str] = "INSTANCE", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), spatial_dims: int = 3, use_downsample: bool = True, node_a=None, @@ -398,7 +474,9 @@ def __init__( bias=False, dilation=1, ), - get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=self.filter_nums[res_idx - 1]), + get_norm_layer( + name=norm_name, spatial_dims=spatial_dims, channels=self.filter_nums[max(res_idx - 1, 0)] + ), nn.Upsample(scale_factor=2 ** (res_idx != 0), mode=mode, align_corners=True), ) @@ -484,6 +562,8 @@ def __init__( num_blocks: int = 6, num_depths: int = 3, spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), use_downsample: bool = True, device: str = "cpu", ): @@ -494,6 +574,8 @@ def __init__( self.num_blocks = num_blocks self.num_depths = num_depths self._spatial_dims = spatial_dims + self._act_name = act_name + self._norm_name = norm_name self.use_downsample = use_downsample self.device = device self.num_cell_ops = 0 @@ -535,6 +617,8 @@ def __init__( self.arch_code2ops[res_idx], self.arch_code_c[blk_idx, res_idx], self._spatial_dims, + self._act_name, + self._norm_name, ) def forward(self, x): @@ -555,6 +639,8 @@ def __init__( num_blocks: int = 6, num_depths: int = 3, spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), use_downsample: bool = True, device: str = "cpu", ): @@ -571,6 +657,8 @@ def __init__( num_blocks=num_blocks, num_depths=num_depths, spatial_dims=spatial_dims, + act_name=act_name, + norm_name=norm_name, use_downsample=use_downsample, device=device, ) @@ -591,7 +679,7 @@ def forward(self, x: List[torch.Tensor]) -> List[torch.Tensor]: x=inputs[self.arch_code2in[res_idx]], weight=torch.ones_like(self.arch_code_c[blk_idx, res_idx]) ) outputs[self.arch_code2out[res_idx]] = outputs[self.arch_code2out[res_idx]] + _out - inputs = outputs + inputs = outputs return inputs @@ -650,6 +738,8 @@ def __init__( num_blocks: int = 6, num_depths: int = 3, spatial_dims: int = 3, + act_name: Union[Tuple, str] = "RELU", + norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}), use_downsample: bool = True, device: str = "cpu", ): @@ -663,6 +753,8 @@ def __init__( num_blocks=num_blocks, num_depths=num_depths, spatial_dims=spatial_dims, + act_name=act_name, + norm_name=norm_name, use_downsample=use_downsample, device=device, ) diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py index 337a99acd8..053ab255b8 100644 --- a/monai/networks/nets/dynunet.py +++ b/monai/networks/nets/dynunet.py @@ -74,10 +74,14 @@ class DynUNet(nn.Module): is no less than 3 in order to have at least one downsample and upsample blocks. To meet the requirements of the structure, the input size for each spatial dimension should be divisible - by `2 * the product of all strides in the corresponding dimension`. The output size for each spatial dimension - equals to the input size of the corresponding dimension divided by the stride in strides[0]. - For example, if `strides=((1, 2, 4), 2, 1, 1)`, the minimal spatial size of the input is `(8, 16, 32)`, and - the spatial size of the output is `(8, 8, 8)`. + by the product of all strides in the corresponding dimension. In addition, the minimal spatial size should have + at least one dimension that has twice the size of the product of all strides. + For example, if `strides=((1, 2, 4), 2, 2, 1)`, the spatial size should be divisible by `(4, 8, 16)`, + and the minimal spatial size is `(8, 8, 16)` or `(4, 16, 16)` or `(4, 8, 32)`. + + The output size for each spatial dimension equals to the input size of the corresponding dimension divided by the + stride in strides[0]. + For example, if `strides=((1, 2, 4), 2, 2, 1)` and the input size is `(64, 32, 32)`, the output size is `(64, 16, 8)`. For backwards compatibility with old weights, please set `strict=False` when calling `load_state_dict`. @@ -100,6 +104,8 @@ class DynUNet(nn.Module): If not specified, the way which nnUNet used will be employed. Defaults to ``None``. dropout: dropout ratio. Defaults to no dropout. norm_name: feature normalization type and arguments. Defaults to ``INSTANCE``. + `INSTANCE_NVFUSER` is a faster version of the instance norm layer, it can be used when: + 1) `spatial_dims=3`, 2) CUDA device is available, 3) `apex` is installed and 4) non-Windows OS is used. act_name: activation layer type and arguments. Defaults to ``leakyrelu``. deep_supervision: whether to add deep supervision head before output. Defaults to ``False``. If ``True``, in training mode, the forward function will output not only the final feature map diff --git a/monai/networks/nets/hovernet.py b/monai/networks/nets/hovernet.py new file mode 100644 index 0000000000..568d6658c5 --- /dev/null +++ b/monai/networks/nets/hovernet.py @@ -0,0 +1,532 @@ +# 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. + +# ========================================================================= +# Adapted from https://github.com/vqdang/hover_net +# which has the following license: +# https://github.com/vqdang/hover_net/blob/master/LICENSE +# MIT License + +# Origial publication: +# @article{graham2019hover, +# title={Hover-net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images}, +# author={Graham, Simon and Vu, Quoc Dang and Raza, Shan E Ahmed and Azam, Ayesha and Tsang, Yee Wah and Kwak, +# Jin Tae and Rajpoot, Nasir}, +# journal={Medical Image Analysis}, +# pages={101563}, +# year={2019}, +# publisher={Elsevier} +# } + +# ========================================================================= + +from collections import OrderedDict +from enum import Enum +from typing import Callable, Dict, List, Sequence, Type, Union + +import torch +import torch.nn as nn + +from monai.networks.blocks import UpSample +from monai.networks.layers.factories import Conv, Dropout +from monai.networks.layers.utils import get_act_layer, get_norm_layer +from monai.utils import InterpolateMode, UpsampleMode, export + +__all__ = ["HoverNet", "Hovernet", "HoVernet", "HoVerNet"] + + +class _DenseLayerDecoder(nn.Module): + def __init__( + self, + num_features: int, + in_channels: int, + out_channels: int, + dropout_prob: float = 0.0, + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + kernel_size: int = 3, + ) -> None: + """ + Args: + spatial_dims: number of spatial dimensions of the input image. + num_features: number of internal channels used for the layer + in_channels: number of the input channels. + out_channels: number of the output channels. + dropout_prob: dropout rate after each dense layer. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + kernel_size: size of the kernel for >1 convolutions (dependent on mode) + """ + super().__init__() + + conv_type: Callable = Conv[Conv.CONV, 2] + dropout_type: Callable = Dropout[Dropout.DROPOUT, 2] + + self.layers = nn.Sequential() + + self.layers.add_module("preact_bna/bn", get_norm_layer(name=norm, spatial_dims=2, channels=in_channels)) + self.layers.add_module("preact_bna/relu", get_act_layer(name=act)) + self.layers.add_module("conv1", conv_type(in_channels, num_features, kernel_size=1, bias=False)) + self.layers.add_module("conv1/norm", get_norm_layer(name=norm, spatial_dims=2, channels=num_features)) + self.layers.add_module("conv1/relu2", get_act_layer(name=act)) + self.layers.add_module( + "conv2", conv_type(num_features, out_channels, kernel_size=kernel_size, padding=0, groups=4, bias=False) + ) + + if dropout_prob > 0: + self.layers.add_module("dropout", dropout_type(dropout_prob)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + x1 = self.layers(x) + if x1.shape != x.shape: + trim = (x.shape[-1] - x1.shape[-1]) // 2 + x = x[:, :, trim:-trim, trim:-trim] + + x = torch.cat([x, x1], 1) + + return x + + +class _DecoderBlock(nn.Sequential): + def __init__( + self, + layers: int, + num_features: int, + in_channels: int, + out_channels: int, + dropout_prob: float = 0.0, + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + kernel_size: int = 3, + ) -> None: + """ + Args: + spatial_dims: number of spatial dimensions of the input image. + layers: number of layers in the block. + num_features: number of internal features used. + in_channels: number of the input channel. + out_channels: number of the output channel. + dropout_prob: dropout rate after each dense layer. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + kernel_size: size of the kernel for >1 convolutions (dependent on mode) + """ + super().__init__() + + conv_type: Callable = Conv[Conv.CONV, 2] + + self.add_module("conva", conv_type(in_channels, in_channels // 4, kernel_size=kernel_size, bias=False)) + + _in_channels = in_channels // 4 + for i in range(layers): + layer = _DenseLayerDecoder( + num_features, _in_channels, out_channels, dropout_prob, act=act, norm=norm, kernel_size=kernel_size + ) + _in_channels += out_channels + self.add_module("denselayerdecoder%d" % (i + 1), layer) + + trans = _Transition(_in_channels, act=act, norm=norm) + self.add_module("bna_block", trans) + self.add_module("convf", conv_type(_in_channels, _in_channels, kernel_size=1, bias=False)) + + +class _DenseLayer(nn.Sequential): + def __init__( + self, + num_features: int, + in_channels: int, + out_channels: int, + dropout_prob: float = 0.0, + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + drop_first_norm_relu: int = 0, + kernel_size: int = 3, + ) -> None: + """Dense Convolutional Block. + + References: + Huang, Gao, et al. "Densely connected convolutional networks." + Proceedings of the IEEE conference on computer vision and + pattern recognition. 2017. + + Args: + num_features: number of internal channels used for the layer + in_channels: number of the input channels. + out_channels: number of the output channels. + dropout_prob: dropout rate after each dense layer. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + drop_first_norm_relu - omits the first norm/relu for the first layer + kernel_size: size of the kernel for >1 convolutions (dependent on mode) + """ + super().__init__() + + self.layers = nn.Sequential() + conv_type: Callable = Conv[Conv.CONV, 2] + dropout_type: Callable = Dropout[Dropout.DROPOUT, 2] + + if not drop_first_norm_relu: + self.layers.add_module("preact_norm", get_norm_layer(name=norm, spatial_dims=2, channels=in_channels)) + self.layers.add_module("preact_relu", get_act_layer(name=act)) + + self.layers.add_module("conv1", conv_type(in_channels, num_features, kernel_size=1, padding=0, bias=False)) + self.layers.add_module("norm2", get_norm_layer(name=norm, spatial_dims=2, channels=num_features)) + self.layers.add_module("relu2", get_act_layer(name=act)) + + if in_channels != 64 and drop_first_norm_relu: + self.layers.add_module( + "conv2", conv_type(num_features, num_features, kernel_size=kernel_size, stride=2, padding=2, bias=False) + ) + else: + self.layers.add_module("conv2", conv_type(num_features, num_features, kernel_size=1, padding=0, bias=False)) + + self.layers.add_module("norm3", get_norm_layer(name=norm, spatial_dims=2, channels=num_features)) + self.layers.add_module("relu3", get_act_layer(name=act)) + self.layers.add_module("conv3", conv_type(num_features, out_channels, kernel_size=1, padding=0, bias=False)) + + if dropout_prob > 0: + self.layers.add_module("dropout", dropout_type(dropout_prob)) + + +class _Transition(nn.Sequential): + def __init__( + self, in_channels: int, act: Union[str, tuple] = ("relu", {"inplace": True}), norm: Union[str, tuple] = "batch" + ) -> None: + """ + Args: + in_channels: number of the input channel. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + """ + super().__init__() + + self.add_module("norm", get_norm_layer(name=norm, spatial_dims=2, channels=in_channels)) + self.add_module("relu", get_act_layer(name=act)) + + +class _ResidualBlock(nn.Module): + def __init__( + self, + layers: int, + num_features: int, + in_channels: int, + out_channels: int, + dropout_prob: float = 0.0, + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + ) -> None: + """Residual block. + + References: + He, Kaiming, et al. "Deep residual learning for image + recognition." Proceedings of the IEEE conference on computer + vision and pattern recognition. 2016. + + Args: + layers: number of layers in the block. + num_features: number of internal features used. + in_channels: number of the input channel. + out_channels: number of the output channel. + dropout_prob: dropout rate after each dense layer. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + """ + super().__init__() + + self.layers = nn.Sequential() + conv_type: Callable = Conv[Conv.CONV, 2] + + if in_channels == 64: + self.shortcut = conv_type(in_channels, out_channels, kernel_size=1, bias=False) + else: + self.shortcut = conv_type(in_channels, out_channels, kernel_size=1, stride=2, padding=1, bias=False) + + layer = _DenseLayer( + num_features, in_channels, out_channels, dropout_prob, act=act, norm=norm, drop_first_norm_relu=True + ) + self.layers.add_module("prim_denselayer%d" % (1), layer) + + for i in range(1, layers): + layer = _DenseLayer(num_features, out_channels, out_channels, dropout_prob, act=act, norm=norm) + self.layers.add_module("main_denselayer%d" % (i + 1), layer) + + self.bna_block = _Transition(out_channels, act=act, norm=norm) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + sc = self.shortcut(x) + + if self.shortcut.stride == (2, 2): + sc = sc[:, :, :-1, :-1] + + for layer in self.layers: + x = layer.forward(x) + if x.shape[-2:] != sc.shape[-2:]: + x = x[:, :, :-1, :-1] + + x = x + sc + sc = x + + x = self.bna_block(x) + + return x + + +class _DecoderBranch(nn.ModuleList): + def __init__( + self, + decode_config: Sequence[int] = (8, 4), + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + dropout_prob: float = 0.0, + out_channels: int = 2, + kernel_size: int = 3, + ) -> None: + """ + Args: + decode_config: number of layers for each block. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + dropout_prob: dropout rate after each dense layer. + num_features: number of internal features used. + out_channels: number of the output channel. + kernel_size: size of the kernel for >1 convolutions (dependent on mode) + """ + super().__init__() + conv_type: Callable = Conv[Conv.CONV, 2] + + # decode branches + _in_channels = 1024 + _num_features = 128 + _out_channels = 32 + + self.decoder_blocks = nn.Sequential() + for i, num_layers in enumerate(decode_config): + block = _DecoderBlock( + layers=num_layers, + num_features=_num_features, + in_channels=_in_channels, + out_channels=_out_channels, + dropout_prob=dropout_prob, + act=act, + norm=norm, + kernel_size=kernel_size, + ) + self.decoder_blocks.add_module(f"decoderblock{i + 1}", block) + _in_channels = 512 + + # output layers + self.output_features = nn.Sequential() + _i = len(decode_config) + _pad_size = (kernel_size - 1) // 2 + _seq_block = nn.Sequential( + OrderedDict( + [("conva", conv_type(256, 64, kernel_size=kernel_size, stride=1, bias=False, padding=_pad_size))] + ) + ) + + self.output_features.add_module(f"decoderblock{_i + 1}", _seq_block) + + _seq_block = nn.Sequential( + OrderedDict( + [ + ("norm", get_norm_layer(name=norm, spatial_dims=2, channels=64)), + ("relu", get_act_layer(name=act)), + ("conv", conv_type(64, out_channels, kernel_size=1, stride=1)), + ] + ) + ) + + self.output_features.add_module(f"decoderblock{_i + 2}", _seq_block) + + self.upsample = UpSample( + 2, scale_factor=2, mode=UpsampleMode.NONTRAINABLE, interp_mode=InterpolateMode.BILINEAR, bias=False + ) + + def forward(self, xin: torch.Tensor, short_cuts: List[torch.Tensor]) -> torch.Tensor: + + block_number = len(short_cuts) - 1 + x = xin + short_cuts[block_number] + + for block in self.decoder_blocks: + x = block(x) + x = self.upsample(x) + block_number -= 1 + trim = (short_cuts[block_number].shape[-1] - x.shape[-1]) // 2 + x += short_cuts[block_number][:, :, trim:-trim, trim:-trim] + + for block in self.output_features: + x = block(x) + + return x + + +@export("monai.networks.nets") +class HoverNet(nn.Module): + """HoVerNet + + References: + Graham, Simon et al. Hover-net: Simultaneous segmentation + and classification of nuclei in multi-tissue histology images, + Medical Image Analysis 2019 + + Args: + in_channels: number of the input channel. + out_classes: number of the nuclear type classes. + act: activation type and arguments. Defaults to relu. + norm: feature normalization type and arguments. Defaults to batch norm. + dropout_prob: dropout rate after each dense layer. + """ + + class Mode(Enum): + FAST: int = 0 + ORIGINAL: int = 1 + + def _mode_to_int(self, mode) -> int: + + if mode == self.Mode.FAST: + return 0 + else: + return 1 + + def __init__( + self, + mode: Mode = Mode.FAST, + in_channels: int = 3, + out_classes: int = 0, + act: Union[str, tuple] = ("relu", {"inplace": True}), + norm: Union[str, tuple] = "batch", + dropout_prob: float = 0.0, + ) -> None: + + super().__init__() + + self.mode: int = self._mode_to_int(mode) + + if mode not in [self.Mode.ORIGINAL, self.Mode.FAST]: + raise ValueError("Input size should be 270 x 270 when using Mode.ORIGINAL") + + if out_classes > 128: + raise ValueError("Number of nuclear types classes exceeds maximum (128)") + elif out_classes == 1: + raise ValueError("Number of nuclear type classes should either be None or >1") + + if dropout_prob > 1 or dropout_prob < 0: + raise ValueError("Dropout can only be in the range 0.0 to 1.0") + + # number of filters in the first convolution layer. + _init_features: int = 64 + # number of layers in each pooling block. + _block_config: Sequence[int] = (3, 4, 6, 3) + + if mode == self.Mode.FAST: + _ksize = 3 + _pad = 3 + else: + _ksize = 5 + _pad = 0 + + conv_type: Type[nn.Conv2d] = Conv[Conv.CONV, 2] + + self.input_features = nn.Sequential( + OrderedDict( + [ + ( + "conv0", + conv_type(in_channels, _init_features, kernel_size=7, stride=1, padding=_pad, bias=False), + ), + ("norm0", get_norm_layer(name=norm, spatial_dims=2, channels=_init_features)), + ("relu0", get_act_layer(name=act)), + ] + ) + ) + + _in_channels = _init_features + _out_channels = 256 + _num_features = _init_features + + self.res_blocks = nn.Sequential() + + for i, num_layers in enumerate(_block_config): + block = _ResidualBlock( + layers=num_layers, + num_features=_num_features, + in_channels=_in_channels, + out_channels=_out_channels, + dropout_prob=dropout_prob, + act=act, + norm=norm, + ) + self.res_blocks.add_module(f"residualblock{i + 1}", block) + + _in_channels = _out_channels + _out_channels *= 2 + _num_features *= 2 + + # bottleneck convolution + self.bottleneck = nn.Sequential() + self.bottleneck.add_module( + "conv_bottleneck", conv_type(_in_channels, _num_features, kernel_size=1, stride=1, padding=0, bias=False) + ) + self.upsample = UpSample( + 2, scale_factor=2, mode=UpsampleMode.NONTRAINABLE, interp_mode=InterpolateMode.BILINEAR, bias=False + ) + + # decode branches + self.nucleus_prediction = _DecoderBranch(kernel_size=_ksize) + self.horizontal_vertical = _DecoderBranch(kernel_size=_ksize) + self.type_prediction: _DecoderBranch = None # type: ignore + + if out_classes > 0: + self.type_prediction = _DecoderBranch(out_channels=out_classes, kernel_size=_ksize) + + for m in self.modules(): + if isinstance(m, conv_type): + nn.init.kaiming_normal_(torch.as_tensor(m.weight)) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(torch.as_tensor(m.weight), 1) + nn.init.constant_(torch.as_tensor(m.bias), 0) + + def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + + if self.mode == 1: + if x.shape[-1] != 270 or x.shape[-2] != 270: + raise ValueError("Input size should be 270 x 270 when using Mode.ORIGINAL") + else: + if x.shape[-1] != 256 or x.shape[-2] != 256: + raise ValueError("Input size should be 256 x 256 when using Mode.FAST") + + x = x / 255.0 # to 0-1 range to match XY + x = self.input_features(x) + short_cuts = [] + + for i, block in enumerate(self.res_blocks): + x = block.forward(x) + + if i <= 2: + short_cuts.append(x) + + x = self.bottleneck(x) + x = self.upsample(x) + + x_np = self.nucleus_prediction(x, short_cuts) + x_hv = self.horizontal_vertical(x, short_cuts) + tp = self.type_prediction + + if tp is not None: + x_tp = self.type_prediction(x, short_cuts) + return {"nucleus_prediction": x_np, "horizonal_vertical": x_hv, "type_prediction": x_tp} + + return {"nucleus_prediction": x_np, "horizonal_vertical": x_hv} + + +Hovernet = HoVernet = HoVerNet = HoverNet diff --git a/monai/networks/nets/swin_unetr.py b/monai/networks/nets/swin_unetr.py new file mode 100644 index 0000000000..994fb50171 --- /dev/null +++ b/monai/networks/nets/swin_unetr.py @@ -0,0 +1,975 @@ +# 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 typing import Sequence, Tuple, Type, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from torch.nn import LayerNorm + +from monai.networks.blocks import MLPBlock as Mlp +from monai.networks.blocks import PatchEmbed, UnetOutBlock, UnetrBasicBlock, UnetrUpBlock +from monai.networks.layers import DropPath, trunc_normal_ +from monai.utils import ensure_tuple_rep, optional_import + +rearrange, _ = optional_import("einops", name="rearrange") + + +class SwinUNETR(nn.Module): + """ + Swin UNETR based on: "Hatamizadeh et al., + Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images + " + """ + + def __init__( + self, + img_size: Union[Sequence[int], int], + in_channels: int, + out_channels: int, + depths: Sequence[int] = (2, 2, 2, 2), + num_heads: Sequence[int] = (3, 6, 12, 24), + feature_size: int = 24, + norm_name: Union[Tuple, str] = "instance", + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + dropout_path_rate: float = 0.0, + normalize: bool = True, + use_checkpoint: bool = False, + spatial_dims: int = 3, + ) -> None: + """ + Args: + img_size: dimension of input image. + in_channels: dimension of input channels. + out_channels: dimension of output channels. + feature_size: dimension of network feature size. + depths: number of layers in each stage. + num_heads: number of attention heads. + norm_name: feature normalization type and arguments. + drop_rate: dropout rate. + attn_drop_rate: attention dropout rate. + dropout_path_rate: drop path rate. + normalize: normalize output intermediate features in each stage. + use_checkpoint: use gradient checkpointing for reduced memory usage. + spatial_dims: number of spatial dims. + + Examples:: + + # for 3D single channel input with size (96,96,96), 4-channel output and feature size of 48. + >>> net = SwinUNETR(img_size=(96,96,96), in_channels=1, out_channels=4, feature_size=48) + + # for 3D 4-channel input with size (128,128,128), 3-channel output and (2,4,2,2) layers in each stage. + >>> net = SwinUNETR(img_size=(128,128,128), in_channels=4, out_channels=3, depths=(2,4,2,2)) + + # for 2D single channel input with size (96,96), 2-channel output and gradient checkpointing. + >>> net = SwinUNETR(img_size=(96,96), in_channels=3, out_channels=2, use_checkpoint=True, spatial_dims=2) + + """ + + super().__init__() + + img_size = ensure_tuple_rep(img_size, spatial_dims) + patch_size = ensure_tuple_rep(2, spatial_dims) + window_size = ensure_tuple_rep(7, spatial_dims) + + if not (spatial_dims == 2 or spatial_dims == 3): + raise ValueError("spatial dimension should be 2 or 3.") + + for m, p in zip(img_size, patch_size): + for i in range(5): + if m % np.power(p, i + 1) != 0: + raise ValueError("input image size (img_size) should be divisible by stage-wise image resolution.") + + if not (0 <= drop_rate <= 1): + raise ValueError("dropout rate should be between 0 and 1.") + + if not (0 <= attn_drop_rate <= 1): + raise ValueError("attention dropout rate should be between 0 and 1.") + + if not (0 <= dropout_path_rate <= 1): + raise ValueError("drop path rate should be between 0 and 1.") + + if feature_size % 12 != 0: + raise ValueError("feature_size should be divisible by 12.") + + self.normalize = normalize + + self.swinViT = SwinTransformer( + in_chans=in_channels, + embed_dim=feature_size, + window_size=window_size, + patch_size=patch_size, + depths=depths, + num_heads=num_heads, + mlp_ratio=4.0, + qkv_bias=True, + drop_rate=drop_rate, + attn_drop_rate=attn_drop_rate, + drop_path_rate=dropout_path_rate, + norm_layer=nn.LayerNorm, + use_checkpoint=use_checkpoint, + spatial_dims=spatial_dims, + ) + + self.encoder1 = UnetrBasicBlock( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=feature_size, + kernel_size=3, + stride=1, + norm_name=norm_name, + res_block=True, + ) + + self.encoder2 = UnetrBasicBlock( + spatial_dims=spatial_dims, + in_channels=feature_size, + out_channels=feature_size, + kernel_size=3, + stride=1, + norm_name=norm_name, + res_block=True, + ) + + self.encoder3 = UnetrBasicBlock( + spatial_dims=spatial_dims, + in_channels=2 * feature_size, + out_channels=2 * feature_size, + kernel_size=3, + stride=1, + norm_name=norm_name, + res_block=True, + ) + + self.encoder4 = UnetrBasicBlock( + spatial_dims=spatial_dims, + in_channels=4 * feature_size, + out_channels=4 * feature_size, + kernel_size=3, + stride=1, + norm_name=norm_name, + res_block=True, + ) + + self.encoder10 = UnetrBasicBlock( + spatial_dims=spatial_dims, + in_channels=16 * feature_size, + out_channels=16 * feature_size, + kernel_size=3, + stride=1, + norm_name=norm_name, + res_block=True, + ) + + self.decoder5 = UnetrUpBlock( + spatial_dims=spatial_dims, + in_channels=16 * feature_size, + out_channels=8 * feature_size, + kernel_size=3, + upsample_kernel_size=2, + norm_name=norm_name, + res_block=True, + ) + + self.decoder4 = UnetrUpBlock( + spatial_dims=spatial_dims, + in_channels=feature_size * 8, + out_channels=feature_size * 4, + kernel_size=3, + upsample_kernel_size=2, + norm_name=norm_name, + res_block=True, + ) + + self.decoder3 = UnetrUpBlock( + spatial_dims=spatial_dims, + in_channels=feature_size * 4, + out_channels=feature_size * 2, + kernel_size=3, + upsample_kernel_size=2, + norm_name=norm_name, + res_block=True, + ) + self.decoder2 = UnetrUpBlock( + spatial_dims=spatial_dims, + in_channels=feature_size * 2, + out_channels=feature_size, + kernel_size=3, + upsample_kernel_size=2, + norm_name=norm_name, + res_block=True, + ) + + self.decoder1 = UnetrUpBlock( + spatial_dims=spatial_dims, + in_channels=feature_size, + out_channels=feature_size, + kernel_size=3, + upsample_kernel_size=2, + norm_name=norm_name, + res_block=True, + ) + + self.out = UnetOutBlock(spatial_dims=spatial_dims, in_channels=feature_size, out_channels=out_channels) + + def load_from(self, weights): + + with torch.no_grad(): + self.swinViT.patch_embed.proj.weight.copy_(weights["state_dict"]["module.patch_embed.proj.weight"]) + self.swinViT.patch_embed.proj.bias.copy_(weights["state_dict"]["module.patch_embed.proj.bias"]) + for bname, block in self.swinViT.layers1[0].blocks.named_children(): + block.load_from(weights, n_block=bname, layer="layers1") + self.swinViT.layers1[0].downsample.reduction.weight.copy_( + weights["state_dict"]["module.layers1.0.downsample.reduction.weight"] + ) + self.swinViT.layers1[0].downsample.norm.weight.copy_( + weights["state_dict"]["module.layers1.0.downsample.norm.weight"] + ) + self.swinViT.layers1[0].downsample.norm.bias.copy_( + weights["state_dict"]["module.layers1.0.downsample.norm.bias"] + ) + for bname, block in self.swinViT.layers2[0].blocks.named_children(): + block.load_from(weights, n_block=bname, layer="layers2") + self.swinViT.layers2[0].downsample.reduction.weight.copy_( + weights["state_dict"]["module.layers2.0.downsample.reduction.weight"] + ) + self.swinViT.layers2[0].downsample.norm.weight.copy_( + weights["state_dict"]["module.layers2.0.downsample.norm.weight"] + ) + self.swinViT.layers2[0].downsample.norm.bias.copy_( + weights["state_dict"]["module.layers2.0.downsample.norm.bias"] + ) + for bname, block in self.swinViT.layers3[0].blocks.named_children(): + block.load_from(weights, n_block=bname, layer="layers3") + self.swinViT.layers3[0].downsample.reduction.weight.copy_( + weights["state_dict"]["module.layers3.0.downsample.reduction.weight"] + ) + self.swinViT.layers3[0].downsample.norm.weight.copy_( + weights["state_dict"]["module.layers3.0.downsample.norm.weight"] + ) + self.swinViT.layers3[0].downsample.norm.bias.copy_( + weights["state_dict"]["module.layers3.0.downsample.norm.bias"] + ) + for bname, block in self.swinViT.layers4[0].blocks.named_children(): + block.load_from(weights, n_block=bname, layer="layers4") + self.swinViT.layers4[0].downsample.reduction.weight.copy_( + weights["state_dict"]["module.layers4.0.downsample.reduction.weight"] + ) + self.swinViT.layers4[0].downsample.norm.weight.copy_( + weights["state_dict"]["module.layers4.0.downsample.norm.weight"] + ) + self.swinViT.layers4[0].downsample.norm.bias.copy_( + weights["state_dict"]["module.layers4.0.downsample.norm.bias"] + ) + + def forward(self, x_in): + hidden_states_out = self.swinViT(x_in, self.normalize) + enc0 = self.encoder1(x_in) + enc1 = self.encoder2(hidden_states_out[0]) + enc2 = self.encoder3(hidden_states_out[1]) + enc3 = self.encoder4(hidden_states_out[2]) + dec4 = self.encoder10(hidden_states_out[4]) + dec3 = self.decoder5(dec4, hidden_states_out[3]) + dec2 = self.decoder4(dec3, enc3) + dec1 = self.decoder3(dec2, enc2) + dec0 = self.decoder2(dec1, enc1) + out = self.decoder1(dec0, enc0) + logits = self.out(out) + return logits + + +def window_partition(x, window_size): + """window partition operation based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + + Args: + x: input tensor. + window_size: local window size. + """ + x_shape = x.size() + if len(x_shape) == 5: + b, d, h, w, c = x_shape + x = x.view( + b, + d // window_size[0], + window_size[0], + h // window_size[1], + window_size[1], + w // window_size[2], + window_size[2], + c, + ) + windows = ( + x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, window_size[0] * window_size[1] * window_size[2], c) + ) + elif len(x_shape) == 4: + b, h, w, c = x.shape + x = x.view(b, h // window_size[0], window_size[0], w // window_size[1], window_size[1], c) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0] * window_size[1], c) + return windows + + +def window_reverse(windows, window_size, dims): + """window reverse operation based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + + Args: + windows: windows tensor. + window_size: local window size. + dims: dimension values. + """ + if len(dims) == 4: + b, d, h, w = dims + x = windows.view( + b, + d // window_size[0], + h // window_size[1], + w // window_size[2], + window_size[0], + window_size[1], + window_size[2], + -1, + ) + x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(b, d, h, w, -1) + + elif len(dims) == 3: + b, h, w = dims + x = windows.view(b, h // window_size[0], w // window_size[0], window_size[0], window_size[1], -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1) + return x + + +def get_window_size(x_size, window_size, shift_size=None): + """Computing window size based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + + Args: + x_size: input size. + window_size: local window size. + shift_size: window shifting size. + """ + + use_window_size = list(window_size) + if shift_size is not None: + use_shift_size = list(shift_size) + for i in range(len(x_size)): + if x_size[i] <= window_size[i]: + use_window_size[i] = x_size[i] + if shift_size is not None: + use_shift_size[i] = 0 + + if shift_size is None: + return tuple(use_window_size) + else: + return tuple(use_window_size), tuple(use_shift_size) + + +class WindowAttention(nn.Module): + """ + Window based multi-head self attention module with relative position bias based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + """ + + def __init__( + self, + dim: int, + num_heads: int, + window_size: Sequence[int], + qkv_bias: bool = False, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + """ + Args: + dim: number of feature channels. + num_heads: number of attention heads. + window_size: local window size. + qkv_bias: add a learnable bias to query, key, value. + attn_drop: attention dropout rate. + proj_drop: dropout rate of output. + """ + + super().__init__() + self.dim = dim + self.window_size = window_size + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + mesh_args = torch.meshgrid.__kwdefaults__ + + if len(self.window_size) == 3: + self.relative_position_bias_table = nn.Parameter( + torch.zeros( + (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1), + num_heads, + ) + ) + coords_d = torch.arange(self.window_size[0]) + coords_h = torch.arange(self.window_size[1]) + coords_w = torch.arange(self.window_size[2]) + if mesh_args is not None: + coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w, indexing="ij")) + else: + coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w)) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += self.window_size[0] - 1 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 2] += self.window_size[2] - 1 + relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1) + relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1 + elif len(self.window_size) == 2: + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + if mesh_args is not None: + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij")) + else: + coords = torch.stack(torch.meshgrid(coords_h, coords_w)) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += self.window_size[0] - 1 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + + relative_position_index = relative_coords.sum(-1) + self.register_buffer("relative_position_index", relative_position_index) + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + trunc_normal_(self.relative_position_bias_table, std=0.02) + self.softmax = nn.Softmax(dim=-1) + + 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) + ].reshape(n, n, -1) + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + attn = attn + relative_position_bias.unsqueeze(0) + if mask is not None: + nw = mask.shape[0] + attn = attn.view(b // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, n, n) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2).reshape(b, n, c) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ + Swin Transformer block based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + """ + + def __init__( + self, + dim: int, + num_heads: int, + window_size: Sequence[int], + shift_size: Sequence[int], + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: str = "GELU", + norm_layer: Type[LayerNorm] = nn.LayerNorm, + use_checkpoint: bool = False, + ) -> None: + """ + Args: + dim: number of feature channels. + num_heads: number of attention heads. + window_size: local window size. + shift_size: window shift size. + mlp_ratio: ratio of mlp hidden dim to embedding dim. + qkv_bias: add a learnable bias to query, key, value. + drop: dropout rate. + attn_drop: attention dropout rate. + drop_path: stochastic depth rate. + act_layer: activation layer. + norm_layer: normalization layer. + use_checkpoint: use gradient checkpointing for reduced memory usage. + """ + + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + self.use_checkpoint = use_checkpoint + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, + window_size=self.window_size, + num_heads=num_heads, + qkv_bias=qkv_bias, + attn_drop=attn_drop, + proj_drop=drop, + ) + + 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 = Mlp(hidden_size=dim, mlp_dim=mlp_hidden_dim, act=act_layer, dropout_rate=drop, dropout_mode="swin") + + def forward_part1(self, x, mask_matrix): + x_shape = x.size() + x = self.norm1(x) + if len(x_shape) == 5: + b, d, h, w, c = x.shape + window_size, shift_size = get_window_size((d, h, w), self.window_size, self.shift_size) + pad_l = pad_t = pad_d0 = 0 + pad_d1 = (window_size[0] - d % window_size[0]) % window_size[0] + pad_b = (window_size[1] - h % window_size[1]) % window_size[1] + pad_r = (window_size[2] - w % window_size[2]) % window_size[2] + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1)) + _, dp, hp, wp, _ = x.shape + dims = [b, dp, hp, wp] + + elif len(x_shape) == 4: + b, h, w, c = x.shape + window_size, shift_size = get_window_size((h, w), self.window_size, self.shift_size) + pad_l = pad_t = 0 + pad_r = (window_size[0] - h % window_size[0]) % window_size[0] + pad_b = (window_size[1] - w % window_size[1]) % window_size[1] + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, hp, wp, _ = x.shape + dims = [b, hp, wp] + + if any(i > 0 for i in shift_size): + if len(x_shape) == 5: + shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3)) + elif len(x_shape) == 4: + shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + x_windows = window_partition(shifted_x, window_size) + attn_windows = self.attn(x_windows, mask=attn_mask) + attn_windows = attn_windows.view(-1, *(window_size + (c,))) + shifted_x = window_reverse(attn_windows, window_size, dims) + if any(i > 0 for i in shift_size): + if len(x_shape) == 5: + x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3)) + elif len(x_shape) == 4: + x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2)) + else: + x = shifted_x + + if len(x_shape) == 5: + if pad_d1 > 0 or pad_r > 0 or pad_b > 0: + x = x[:, :d, :h, :w, :].contiguous() + elif len(x_shape) == 4: + if pad_r > 0 or pad_b > 0: + x = x[:, :h, :w, :].contiguous() + + return x + + def forward_part2(self, x): + return self.drop_path(self.mlp(self.norm2(x))) + + def load_from(self, weights, n_block, layer): + root = f"module.{layer}.0.blocks.{n_block}." + block_names = [ + "norm1.weight", + "norm1.bias", + "attn.relative_position_bias_table", + "attn.relative_position_index", + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "norm2.weight", + "norm2.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + ] + with torch.no_grad(): + self.norm1.weight.copy_(weights["state_dict"][root + block_names[0]]) + self.norm1.bias.copy_(weights["state_dict"][root + block_names[1]]) + self.attn.relative_position_bias_table.copy_(weights["state_dict"][root + block_names[2]]) + self.attn.relative_position_index.copy_(weights["state_dict"][root + block_names[3]]) + self.attn.qkv.weight.copy_(weights["state_dict"][root + block_names[4]]) + self.attn.qkv.bias.copy_(weights["state_dict"][root + block_names[5]]) + self.attn.proj.weight.copy_(weights["state_dict"][root + block_names[6]]) + self.attn.proj.bias.copy_(weights["state_dict"][root + block_names[7]]) + self.norm2.weight.copy_(weights["state_dict"][root + block_names[8]]) + self.norm2.bias.copy_(weights["state_dict"][root + block_names[9]]) + self.mlp.linear1.weight.copy_(weights["state_dict"][root + block_names[10]]) + self.mlp.linear1.bias.copy_(weights["state_dict"][root + block_names[11]]) + self.mlp.linear2.weight.copy_(weights["state_dict"][root + block_names[12]]) + self.mlp.linear2.bias.copy_(weights["state_dict"][root + block_names[13]]) + + def forward(self, x, mask_matrix): + shortcut = x + if self.use_checkpoint: + x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix) + else: + x = self.forward_part1(x, mask_matrix) + x = shortcut + self.drop_path(x) + if self.use_checkpoint: + x = x + checkpoint.checkpoint(self.forward_part2, x) + else: + x = x + self.forward_part2(x) + return x + + +class PatchMerging(nn.Module): + """ + Patch merging layer based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + """ + + def __init__(self, dim: int, norm_layer: Type[LayerNorm] = nn.LayerNorm, spatial_dims: int = 3) -> None: + """ + Args: + dim: number of feature channels. + norm_layer: normalization layer. + spatial_dims: number of spatial dims. + """ + + super().__init__() + self.dim = dim + if spatial_dims == 3: + self.reduction = nn.Linear(8 * dim, 2 * dim, bias=False) + self.norm = norm_layer(8 * dim) + elif spatial_dims == 2: + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x): + + x_shape = x.size() + if len(x_shape) == 5: + b, d, h, w, c = x_shape + pad_input = (h % 2 == 1) or (w % 2 == 1) or (d % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, d % 2, 0, w % 2, 0, h % 2)) + x0 = x[:, 0::2, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, 0::2, :] + x2 = x[:, 0::2, 1::2, 0::2, :] + x3 = x[:, 0::2, 0::2, 1::2, :] + x4 = x[:, 1::2, 0::2, 1::2, :] + x5 = x[:, 0::2, 1::2, 0::2, :] + x6 = x[:, 0::2, 0::2, 1::2, :] + x7 = x[:, 1::2, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1) + + elif len(x_shape) == 4: + b, h, w, c = x_shape + pad_input = (h % 2 == 1) or (w % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, w % 2, 0, h % 2)) + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], -1) + + x = self.norm(x) + x = self.reduction(x) + return x + + +def compute_mask(dims, window_size, shift_size, device): + """Computing region masks based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + + Args: + dims: dimension values. + window_size: local window size. + shift_size: shift size. + device: device. + """ + + cnt = 0 + + if len(dims) == 3: + d, h, w = dims + img_mask = torch.zeros((1, d, h, w, 1), device=device) + for d in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0], None): + for h in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1], None): + for w in slice(-window_size[2]), slice(-window_size[2], -shift_size[2]), slice(-shift_size[2], None): + img_mask[:, d, h, w, :] = cnt + cnt += 1 + + elif len(dims) == 2: + h, w = dims + img_mask = torch.zeros((1, h, w, 1), device=device) + for h in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0], None): + for w in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1], None): + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, window_size) + mask_windows = mask_windows.squeeze(-1) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + + return attn_mask + + +class BasicLayer(nn.Module): + """ + Basic Swin Transformer layer in one stage based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + """ + + def __init__( + self, + dim: int, + depth: int, + num_heads: int, + window_size: Sequence[int], + drop_path: list, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + drop: float = 0.0, + attn_drop: float = 0.0, + norm_layer: Type[LayerNorm] = nn.LayerNorm, + downsample: isinstance = None, # type: ignore + use_checkpoint: bool = False, + ) -> None: + """ + Args: + dim: number of feature channels. + depths: number of layers in each stage. + num_heads: number of attention heads. + window_size: local window size. + drop_path: stochastic depth rate. + mlp_ratio: ratio of mlp hidden dim to embedding dim. + qkv_bias: add a learnable bias to query, key, value. + drop: dropout rate. + attn_drop: attention dropout rate. + norm_layer: normalization layer. + downsample: downsample layer at the end of the layer. + use_checkpoint: use gradient checkpointing for reduced memory usage. + """ + + super().__init__() + self.window_size = window_size + self.shift_size = tuple(i // 2 for i in window_size) + 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.downsample = downsample + if self.downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer, spatial_dims=len(self.window_size)) + + def forward(self, x): + x_shape = x.size() + if len(x_shape) == 5: + b, c, d, h, w = x_shape + window_size, shift_size = get_window_size((d, h, w), self.window_size, self.shift_size) + x = rearrange(x, "b c d h w -> b d h w c") + 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) + for blk in self.blocks: + x = blk(x, attn_mask) + x = x.view(b, d, h, w, -1) + if self.downsample is not None: + x = self.downsample(x) + x = rearrange(x, "b d h w c -> b c d h w") + + elif len(x_shape) == 4: + b, c, h, w = x_shape + window_size, shift_size = get_window_size((h, w), self.window_size, self.shift_size) + 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) + for blk in self.blocks: + x = blk(x, attn_mask) + x = x.view(b, h, w, -1) + if self.downsample is not None: + x = self.downsample(x) + x = rearrange(x, "b h w c -> b c h w") + return x + + +class SwinTransformer(nn.Module): + """ + Swin Transformer based on: "Liu et al., + Swin Transformer: Hierarchical Vision Transformer using Shifted Windows + " + https://github.com/microsoft/Swin-Transformer + """ + + def __init__( + self, + in_chans: int, + embed_dim: int, + window_size: Sequence[int], + patch_size: Sequence[int], + depths: Sequence[int], + num_heads: Sequence[int], + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + norm_layer: Type[LayerNorm] = nn.LayerNorm, + patch_norm: bool = False, + use_checkpoint: bool = False, + spatial_dims: int = 3, + ) -> None: + """ + Args: + in_chans: dimension of input channels. + embed_dim: number of linear projection output channels. + window_size: local window size. + patch_size: patch size. + depths: number of layers in each stage. + num_heads: number of attention heads. + mlp_ratio: ratio of mlp hidden dim to embedding dim. + qkv_bias: add a learnable bias to query, key, value. + drop_rate: dropout rate. + attn_drop_rate: attention dropout rate. + drop_path_rate: stochastic depth rate. + norm_layer: normalization layer. + patch_norm: add normalization after patch embedding. + use_checkpoint: use gradient checkpointing for reduced memory usage. + spatial_dims: spatial dimension. + """ + + super().__init__() + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.patch_norm = patch_norm + self.window_size = window_size + self.patch_size = patch_size + self.patch_embed = PatchEmbed( + patch_size=self.patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, # type: ignore + spatial_dims=spatial_dims, + ) + self.pos_drop = nn.Dropout(p=drop_rate) + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + self.layers1 = nn.ModuleList() + self.layers2 = nn.ModuleList() + self.layers3 = nn.ModuleList() + self.layers4 = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=self.window_size, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop_rate, + attn_drop=attn_drop_rate, + norm_layer=norm_layer, + downsample=PatchMerging, + use_checkpoint=use_checkpoint, + ) + if i_layer == 0: + self.layers1.append(layer) + elif i_layer == 1: + self.layers2.append(layer) + elif i_layer == 2: + self.layers3.append(layer) + elif i_layer == 3: + self.layers4.append(layer) + self.num_features = int(embed_dim * 2 ** (self.num_layers - 1)) + + def proj_out(self, x, normalize=False): + if normalize: + x_shape = x.size() + if len(x_shape) == 5: + n, ch, d, h, w = x_shape + x = rearrange(x, "n c d h w -> n d h w c") + x = F.layer_norm(x, [ch]) + x = rearrange(x, "n d h w c -> n c d h w") + elif len(x_shape) == 4: + n, ch, h, w = x_shape + x = rearrange(x, "n c h w -> n h w c") + x = F.layer_norm(x, [ch]) + x = rearrange(x, "n h w c -> n c h w") + return x + + def forward(self, x, normalize=True): + x0 = self.patch_embed(x) + x0 = self.pos_drop(x0) + x0_out = self.proj_out(x0, normalize) + x1 = self.layers1[0](x0.contiguous()) + x1_out = self.proj_out(x1, normalize) + x2 = self.layers2[0](x1.contiguous()) + x2_out = self.proj_out(x2, normalize) + x3 = self.layers3[0](x2.contiguous()) + x3_out = self.proj_out(x3, normalize) + x4 = self.layers4[0](x3.contiguous()) + x4_out = self.proj_out(x4, normalize) + return [x0_out, x1_out, x2_out, x3_out, x4_out] diff --git a/monai/networks/nets/torchvision_fc.py b/monai/networks/nets/torchvision_fc.py index e93019d050..ddfee6e041 100644 --- a/monai/networks/nets/torchvision_fc.py +++ b/monai/networks/nets/torchvision_fc.py @@ -9,15 +9,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple from monai.networks.nets import NetAdapter -from monai.utils import deprecated, deprecated_arg, optional_import +from monai.utils import deprecated_arg, optional_import models, _ = optional_import("torchvision.models") -__all__ = ["TorchVisionFCModel", "TorchVisionFullyConvModel"] +__all__ = ["TorchVisionFCModel"] class TorchVisionFCModel(NetAdapter): @@ -72,44 +72,3 @@ def __init__( pool=pool, bias=bias, ) - - -@deprecated(since="0.6.0", removed="0.9.0", msg_suffix="Please consider using `TorchVisionFCModel` instead.") -class TorchVisionFullyConvModel(TorchVisionFCModel): - """ - Customize TorchVision models to replace fully connected layer by convolutional layer. - - Args: - model_name: name of any torchvision with adaptive avg pooling and fully connected layer at the end. - ``resnet18`` (default), ``resnet34``, ``resnet50``, ``resnet101``, ``resnet152``, - ``resnext50_32x4d``, ``resnext101_32x8d``, ``wide_resnet50_2``, ``wide_resnet101_2``. - num_classes: number of classes for the last classification layer. Default to 1. - pool_size: the kernel size for `AvgPool2d` to replace `AdaptiveAvgPool2d`. Default to (7, 7). - pool_stride: the stride for `AvgPool2d` to replace `AdaptiveAvgPool2d`. Default to 1. - pretrained: whether to use the imagenet pretrained weights. Default to False. - - .. deprecated:: 0.6.0 - Use :class:`monai.networks.nets.TorchVisionFCModel` instead. - - """ - - @deprecated_arg("n_classes", since="0.6") - def __init__( - self, - model_name: str = "resnet18", - num_classes: int = 1, - pool_size: Union[int, Tuple[int, int]] = (7, 7), - pool_stride: Union[int, Tuple[int, int]] = 1, - pretrained: bool = False, - n_classes: Optional[int] = None, - ): - # in case the new num_classes is default but you still call deprecated n_classes - if n_classes is not None and num_classes == 1: - num_classes = n_classes - super().__init__( - model_name=model_name, - num_classes=num_classes, - use_conv=True, - pool=("avg", {"kernel_size": pool_size, "stride": pool_stride}), - pretrained=pretrained, - ) diff --git a/monai/networks/nets/unet.py b/monai/networks/nets/unet.py index 25ce61ab3a..faccddee45 100644 --- a/monai/networks/nets/unet.py +++ b/monai/networks/nets/unet.py @@ -68,6 +68,8 @@ class UNet(nn.Module): bias: whether to have a bias term in convolution blocks. Defaults to True. According to `Performance Tuning Guide `_, if a conv layer is directly followed by a batch norm layer, bias should be False. + adn_ordering: a string representing the ordering of activation (A), normalization (N), and dropout (D). + Defaults to "NDA". See also: :py:class:`monai.networks.blocks.ADN`. Examples:: @@ -122,6 +124,7 @@ def __init__( norm: Union[Tuple, str] = Norm.INSTANCE, dropout: float = 0.0, bias: bool = True, + adn_ordering: str = "NDA", dimensions: Optional[int] = None, ) -> None: @@ -155,6 +158,7 @@ def __init__( self.norm = norm self.dropout = dropout self.bias = bias + self.adn_ordering = adn_ordering def _create_block( inc: int, outc: int, channels: Sequence[int], strides: Sequence[int], is_top: bool @@ -229,6 +233,7 @@ def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_ norm=self.norm, dropout=self.dropout, bias=self.bias, + adn_ordering=self.adn_ordering, ) return mod mod = Convolution( @@ -241,6 +246,7 @@ def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_ norm=self.norm, dropout=self.dropout, bias=self.bias, + adn_ordering=self.adn_ordering, ) return mod @@ -279,6 +285,7 @@ def _get_up_layer(self, in_channels: int, out_channels: int, strides: int, is_to bias=self.bias, conv_only=is_top and self.num_res_units == 0, is_transposed=True, + adn_ordering=self.adn_ordering, ) if self.num_res_units > 0: @@ -294,6 +301,7 @@ def _get_up_layer(self, in_channels: int, out_channels: int, strides: int, is_to dropout=self.dropout, bias=self.bias, last_conv_only=is_top, + adn_ordering=self.adn_ordering, ) conv = nn.Sequential(conv, ru) diff --git a/monai/networks/utils.py b/monai/networks/utils.py index a6b0699107..f9018b0c39 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -15,7 +15,8 @@ import warnings from collections import OrderedDict from contextlib import contextmanager -from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union +from copy import deepcopy +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union import torch import torch.nn as nn @@ -23,7 +24,6 @@ from monai.config import PathLike from monai.utils.deprecate_utils import deprecated, deprecated_arg from monai.utils.misc import ensure_tuple, save_obj, set_determinism -from monai.utils.module import pytorch_after __all__ = [ "one_hot", @@ -41,6 +41,8 @@ "save_state", "convert_to_torchscript", "meshgrid_ij", + "replace_modules", + "replace_modules_temp", ] @@ -135,11 +137,17 @@ def normalize_transform( device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, align_corners: bool = False, + zero_centered: bool = False, ) -> torch.Tensor: """ Compute an affine matrix according to the input shape. The transform normalizes the homogeneous image coordinates to the - range of `[-1, 1]`. + range of `[-1, 1]`. Currently the following source coordinates are supported: + + - `align_corners=False`, `zero_centered=False`, normalizing from ``[-0.5, d-0.5]``. + - `align_corners=True`, `zero_centered=False`, normalizing from ``[0, d-1]``. + - `align_corners=False`, `zero_centered=True`, normalizing from ``[-(d+1)/2, (d-1)/2]``. + - `align_corners=True`, `zero_centered=True`, normalizing from ``[-(d-1)/2, (d-1)/2]``. Args: shape: input spatial shape @@ -148,25 +156,32 @@ def normalize_transform( align_corners: if True, consider -1 and 1 to refer to the centers of the corner pixels rather than the image corners. See also: https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample + zero_centered: whether the coordinates are normalized from a zero-centered range, default to `False`. + Setting this flag and `align_corners` will jointly specify the normalization source range. """ norm = torch.tensor(shape, dtype=torch.float64, device=device) # no in-place change if align_corners: norm[norm <= 1.0] = 2.0 norm = 2.0 / (norm - 1.0) norm = torch.diag(torch.cat((norm, torch.ones((1,), dtype=torch.float64, device=device)))) - norm[:-1, -1] = -1.0 + if not zero_centered: # else shift is 0 + norm[:-1, -1] = -1.0 else: norm[norm <= 0.0] = 2.0 norm = 2.0 / norm norm = torch.diag(torch.cat((norm, torch.ones((1,), dtype=torch.float64, device=device)))) - norm[:-1, -1] = 1.0 / torch.tensor(shape, dtype=torch.float64, device=device) - 1.0 + norm[:-1, -1] = 1.0 / torch.tensor(shape, dtype=torch.float64, device=device) - (0.0 if zero_centered else 1.0) norm = norm.unsqueeze(0).to(dtype=dtype) norm.requires_grad = False return norm def to_norm_affine( - affine: torch.Tensor, src_size: Sequence[int], dst_size: Sequence[int], align_corners: bool = False + affine: torch.Tensor, + src_size: Sequence[int], + dst_size: Sequence[int], + align_corners: bool = False, + zero_centered: bool = False, ) -> torch.Tensor: """ Given ``affine`` defined for coordinates in the pixel space, compute the corresponding affine @@ -179,6 +194,8 @@ def to_norm_affine( align_corners: if True, consider -1 and 1 to refer to the centers of the corner pixels rather than the image corners. See also: https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample + zero_centered: whether the coordinates are normalized from a zero-centered range, default to `False`. + See also: :py:func:`monai.networks.utils.normalize_transform`. Raises: TypeError: When ``affine`` is not a ``torch.Tensor``. @@ -194,8 +211,8 @@ def to_norm_affine( if sr != len(src_size) or sr != len(dst_size): raise ValueError(f"affine suggests {sr}D, got src={len(src_size)}D, dst={len(dst_size)}D.") - src_xform = normalize_transform(src_size, affine.device, affine.dtype, align_corners) - dst_xform = normalize_transform(dst_size, affine.device, affine.dtype, align_corners) + src_xform = normalize_transform(src_size, affine.device, affine.dtype, align_corners, zero_centered) + dst_xform = normalize_transform(dst_size, affine.device, affine.dtype, align_corners, zero_centered) return src_xform @ affine @ torch.inverse(dst_xform) @@ -281,14 +298,16 @@ def pixelshuffle( f"divisible by scale_factor ** dimensions ({factor}**{dim}={scale_divisor})." ) - org_channels = channels // scale_divisor + org_channels = int(channels // scale_divisor) output_size = [batch_size, org_channels] + [d * factor for d in input_size[2:]] - indices = tuple(range(2, 2 + 2 * dim)) - indices_factor, indices_dim = indices[:dim], indices[dim:] - permute_indices = (0, 1) + sum(zip(indices_dim, indices_factor), ()) + indices = list(range(2, 2 + 2 * dim)) + indices = indices[dim:] + indices[:dim] + permute_indices = [0, 1] + for idx in range(dim): + permute_indices.extend(indices[idx::dim]) - x = x.reshape(batch_size, org_channels, *([factor] * dim + input_size[2:])) + x = x.reshape([batch_size, org_channels] + [factor] * dim + input_size[2:]) x = x.permute(permute_indices).reshape(output_size) return x @@ -467,7 +486,7 @@ def save_state(src: Union[torch.nn.Module, Dict], path: PathLike, **kwargs): src: input data to save, can be `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`. path: target file path to save the input object. kwargs: other args for the `save_obj` except for the `obj` and `path`. - default `func` is `torch.save()`, details of the args of it: + default `func` is `torch.save()`, details of the args: https://pytorch.org/docs/stable/generated/torch.save.html. """ @@ -502,7 +521,6 @@ def convert_to_torchscript( filename_or_obj: if not None, specify a file-like object (has to implement write and flush) or a string containing a file path name to save the TorchScript model. extra_files: map from filename to contents which will be stored as part of the save model file. - works for PyTorch 1.7 or later. for more details: https://pytorch.org/docs/stable/generated/torch.jit.save.html. verify: whether to verify the input and output of TorchScript model. if `filename_or_obj` is not None, load the saved TorchScript model and verify. @@ -519,10 +537,7 @@ def convert_to_torchscript( with torch.no_grad(): script_module = torch.jit.script(model, **kwargs) if filename_or_obj is not None: - if not pytorch_after(1, 7): - torch.jit.save(m=script_module, f=filename_or_obj) - else: - torch.jit.save(m=script_module, f=filename_or_obj, _extra_files=extra_files) + torch.jit.save(m=script_module, f=filename_or_obj, _extra_files=extra_files) if verify: if device is None: @@ -550,6 +565,105 @@ def convert_to_torchscript( def meshgrid_ij(*tensors): - if pytorch_after(1, 10): - return torch.meshgrid(*tensors, indexing="ij") + if torch.meshgrid.__kwdefaults__ is not None and "indexing" in torch.meshgrid.__kwdefaults__: + return torch.meshgrid(*tensors, indexing="ij") # new api pytorch after 1.10 return torch.meshgrid(*tensors) + + +def _replace_modules( + parent: torch.nn.Module, + name: str, + new_module: torch.nn.Module, + out: List[Tuple[str, torch.nn.Module]], + strict_match: bool = True, + match_device: bool = True, +) -> None: + """ + Helper function for :py:class:`monai.networks.utils.replace_modules`. + """ + if match_device: + devices = list({i.device for i in parent.parameters()}) + # if only one device for whole of model + if len(devices) == 1: + new_module.to(devices[0]) + idx = name.find(".") + # if there is "." in name, call recursively + if idx != -1: + parent_name = name[:idx] + parent = getattr(parent, parent_name) + name = name[idx + 1 :] + _out: List[Tuple[str, torch.nn.Module]] = [] + _replace_modules(parent, name, new_module, _out) + # prepend the parent name + out += [(f"{parent_name}.{r[0]}", r[1]) for r in _out] + # no "." in module name, do the actual replacing + else: + if strict_match: + old_module = getattr(parent, name) + setattr(parent, name, new_module) + out += [(name, old_module)] + else: + for mod_name, _ in parent.named_modules(): + if name in mod_name: + _replace_modules(parent, mod_name, deepcopy(new_module), out, strict_match=True) + + +def replace_modules( + parent: torch.nn.Module, + name: str, + new_module: torch.nn.Module, + strict_match: bool = True, + match_device: bool = True, +) -> List[Tuple[str, torch.nn.Module]]: + """ + Replace sub-module(s) in a parent module. + + The name of the module to be replace can be nested e.g., + `features.denseblock1.denselayer1.layers.relu1`. If this is the case (there are "." + in the module name), then this function will recursively call itself. + + Args: + parent: module that contains the module to be replaced + name: name of module to be replaced. Can include ".". + new_module: `torch.nn.Module` to be placed at position `name` inside `parent`. This will + be deep copied if `strict_match == False` multiple instances are independent. + strict_match: if `True`, module name must `== name`. If false then + `name in named_modules()` will be used. `True` can be used to change just + one module, whereas `False` can be used to replace all modules with similar + name (e.g., `relu`). + match_device: if `True`, the device of the new module will match the model. Requires all + of `parent` to be on the same device. + + Returns: + List of tuples of replaced modules. Element 0 is module name, element 1 is the replaced module. + + Raises: + AttributeError: if `strict_match` is `True` and `name` is not a named module in `parent`. + """ + out: List[Tuple[str, torch.nn.Module]] = [] + _replace_modules(parent, name, new_module, out, strict_match, match_device) + return out + + +@contextmanager +def replace_modules_temp( + parent: torch.nn.Module, + name: str, + new_module: torch.nn.Module, + strict_match: bool = True, + match_device: bool = True, +): + """ + Temporarily replace sub-module(s) in a parent module (context manager). + + See :py:class:`monai.networks.utils.replace_modules`. + """ + replaced: List[Tuple[str, torch.nn.Module]] = [] + try: + # replace + _replace_modules(parent, name, new_module, replaced, strict_match, match_device) + yield + finally: + # revert + for name, module in replaced: + _replace_modules(parent, name, module, [], strict_match=True, match_device=match_device) diff --git a/monai/optimizers/utils.py b/monai/optimizers/utils.py index 1a040927d8..5444aca191 100644 --- a/monai/optimizers/utils.py +++ b/monai/optimizers/utils.py @@ -37,7 +37,7 @@ def generate_param_groups( for "select", the parameters will be `select_func(network).parameters()`. for "filter", the parameters will be - `map(lambda x: x[1], filter(filter_func, network.named_parameters()))` + `(x[1] for x in filter(f, network.named_parameters()))` match_types: a list of tags to identify the matching type corresponding to the `layer_matches` functions, can be "select" or "filter". lr_values: a list of LR values corresponding to the `layer_matches` functions. @@ -76,7 +76,7 @@ def _select(): def _get_filter(f): def _filter(): # should eventually generate a list of network parameters - return map(lambda x: x[1], filter(f, network.named_parameters())) + return (x[1] for x in filter(f, network.named_parameters())) return _filter @@ -91,7 +91,7 @@ def _filter(): raise ValueError(f"unsupported layer match type: {ty}.") params.append({"params": layer_params(), "lr": lr}) - _layers.extend(list(map(id, layer_params()))) + _layers.extend([id(x) for x in layer_params()]) if include_others: params.append({"params": filter(lambda p: id(p) not in _layers, network.parameters())}) diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index a3dc439a51..0cc7fb3e67 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -16,6 +16,7 @@ BoundingRect, CenterScaleCrop, CenterSpatialCrop, + Crop, CropForeground, DivisiblePad, Pad, @@ -43,19 +44,27 @@ CenterSpatialCropd, CenterSpatialCropD, CenterSpatialCropDict, + Cropd, + CropD, + CropDict, CropForegroundd, CropForegroundD, CropForegroundDict, DivisiblePadd, DivisiblePadD, DivisiblePadDict, - PadModeSequence, + Padd, + PadD, + PadDict, RandCropByLabelClassesd, RandCropByLabelClassesD, RandCropByLabelClassesDict, RandCropByPosNegLabeld, RandCropByPosNegLabelD, RandCropByPosNegLabelDict, + RandCropd, + RandCropD, + RandCropDict, RandScaleCropd, RandScaleCropD, RandScaleCropDict, @@ -81,6 +90,7 @@ from .intensity.array import ( AdjustContrast, DetectEnvelope, + ForegroundMask, GaussianSharpen, GaussianSmooth, GibbsNoise, @@ -117,6 +127,9 @@ AdjustContrastd, AdjustContrastD, AdjustContrastDict, + ForegroundMaskd, + ForegroundMaskD, + ForegroundMaskDict, GaussianSharpend, GaussianSharpenD, GaussianSharpenDict, @@ -206,6 +219,14 @@ from .inverse_batch_transform import BatchInverseTransform, Decollated, DecollateD, DecollateDict from .io.array import SUPPORTED_READERS, LoadImage, SaveImage from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict +from .meta_utility.dictionary import ( + FromMetaTensord, + FromMetaTensorD, + FromMetaTensorDict, + ToMetaTensord, + ToMetaTensorD, + ToMetaTensorDict, +) from .nvtx import ( Mark, Markd, @@ -303,6 +324,8 @@ AffineGrid, Flip, GridDistortion, + GridPatch, + GridSplit, Orientation, Rand2DElastic, Rand3DElastic, @@ -312,6 +335,7 @@ RandDeformGrid, RandFlip, RandGridDistortion, + RandGridPatch, RandRotate, RandRotate90, RandZoom, @@ -334,6 +358,12 @@ GridDistortiond, GridDistortionD, GridDistortionDict, + GridPatchd, + GridPatchD, + GridPatchDict, + GridSplitd, + GridSplitD, + GridSplitDict, Orientationd, OrientationD, OrientationDict, @@ -355,6 +385,9 @@ RandGridDistortiond, RandGridDistortionD, RandGridDistortionDict, + RandGridPatchd, + RandGridPatchD, + RandGridPatchDict, RandRotate90d, RandRotate90D, RandRotate90Dict, @@ -412,6 +445,7 @@ RepeatChannel, SimulateDelay, SplitChannel, + SplitDim, SqueezeDim, ToCupy, ToDevice, @@ -509,6 +543,9 @@ SplitChanneld, SplitChannelD, SplitChannelDict, + SplitDimd, + SplitDimD, + SplitDimDict, SqueezeDimd, SqueezeDimD, SqueezeDimDict, @@ -538,7 +575,7 @@ Fourier, allow_missing_keys_mode, compute_divisible_spatial_size, - convert_inverse_interp_mode, + convert_applied_interp_mode, convert_pad_mode, convert_to_contiguous, copypaste_arrays, diff --git a/monai/transforms/compose.py b/monai/transforms/compose.py index bc55af0b15..1d60c34c3e 100644 --- a/monai/transforms/compose.py +++ b/monai/transforms/compose.py @@ -109,7 +109,7 @@ class Compose(Randomizable, InvertibleTransform): defaults to `False`. log_stats: whether to log the detailed information of data and applied transform when error happened, for NumPy array and PyTorch Tensor, log the data shape and value range, - for other meta data, log the values directly. default to `False`. + for other metadata, log the values directly. default to `False`. """ @@ -199,7 +199,7 @@ class OneOf(Compose): defaults to `False`. log_stats: whether to log the detailed information of data and applied transform when error happened, for NumPy array and PyTorch Tensor, log the data shape and value range, - for other meta data, log the values directly. default to `False`. + for other metadata, log the values directly. default to `False`. """ diff --git a/monai/transforms/croppad/array.py b/monai/transforms/croppad/array.py index 55718117b4..662eb62eb6 100644 --- a/monai/transforms/croppad/array.py +++ b/monai/transforms/croppad/array.py @@ -23,11 +23,15 @@ from monai.config import IndexSelection 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.utils import get_random_patch, get_valid_patch_size +from monai.transforms.inverse import InvertibleTransform, TraceableTransform from monai.transforms.transform import Randomizable, Transform from monai.transforms.utils import ( compute_divisible_spatial_size, convert_pad_mode, + create_translate, generate_label_classes_crop_centers, generate_pos_neg_label_crop_centers, generate_spatial_bounding_box, @@ -36,24 +40,28 @@ map_classes_to_indices, weighted_patch_samples, ) -from monai.transforms.utils_pytorch_numpy_unification import floor_divide, maximum +from monai.utils import ImageMetaKey as Key from monai.utils import ( Method, - NumpyPadMode, PytorchPadMode, + TraceKeys, + TransformBackends, + convert_data_type, + convert_to_dst_type, + convert_to_tensor, ensure_tuple, ensure_tuple_rep, fall_back_tuple, look_up_option, + pytorch_after, ) -from monai.utils.enums import TransformBackends -from monai.utils.type_conversion import convert_data_type, convert_to_dst_type __all__ = [ "Pad", "SpatialPad", "BorderPad", "DivisiblePad", + "Crop", "SpatialCrop", "CenterSpatialCrop", "CenterScaleCrop", @@ -69,13 +77,16 @@ ] -class Pad(Transform): +class Pad(InvertibleTransform): """ Perform padding for a given an amount of padding in each dimension. - If input is `torch.Tensor`, `torch.nn.functional.pad` will be used, otherwise, `np.pad` will be used. + + `torch.nn.functional.pad` is used unless the mode or kwargs are not available in torch, + in which case `np.pad` will be used. Args: to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...]. + if None, must provide in the `__call__` at runtime. mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. @@ -84,66 +95,118 @@ class Pad(Transform): https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html kwargs: other arguments for the `np.pad` or `torch.pad` function. note that `np.pad` treats channel dimension as the first dimension. + """ backend = [TransformBackends.TORCH, TransformBackends.NUMPY] def __init__( - self, - to_pad: List[Tuple[int, int]], - mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT, - **kwargs, + self, to_pad: Optional[List[Tuple[int, int]]] = None, mode: str = PytorchPadMode.CONSTANT, **kwargs ) -> None: self.to_pad = to_pad self.mode = mode self.kwargs = kwargs + def compute_pad_width(self, spatial_shape: Sequence[int]) -> List[Tuple[int, int]]: + """ + dynamically compute the pad width according to the spatial shape. + the output is the amount of padding for all dimensions including the channel. + + Args: + spatial_shape: spatial shape of the original image. + + """ + raise NotImplementedError(f"subclass {self.__class__.__name__} must implement this method.") + @staticmethod - def _np_pad(img: np.ndarray, all_pad_width, mode, **kwargs) -> np.ndarray: - return np.pad(img, all_pad_width, mode=mode, **kwargs) # type: ignore + def _np_pad(img: torch.Tensor, pad_width, mode, **kwargs) -> torch.Tensor: + img_np = img.detach().cpu().numpy() if isinstance(img, torch.Tensor) else img + mode = convert_pad_mode(dst=img_np, mode=mode).value + out = torch.as_tensor(np.pad(img, pad_width, mode=mode, **kwargs)) + if isinstance(img, MetaTensor): + out = MetaTensor(out, meta=img.meta, applied_operations=img.applied_operations) + return out @staticmethod - def _pt_pad(img: torch.Tensor, all_pad_width, mode, **kwargs) -> torch.Tensor: - pt_pad_width = [val for sublist in all_pad_width[1:] for val in sublist[::-1]][::-1] + def _pt_pad(img: torch.Tensor, pad_width, mode, **kwargs) -> torch.Tensor: + pt_pad_width = [val for sublist in pad_width[1:] for val in sublist[::-1]][::-1] # torch.pad expects `[B, C, H, W, [D]]` shape return pad_pt(img.unsqueeze(0), pt_pad_width, mode=mode, **kwargs).squeeze(0) - def __call__( - self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None - ) -> NdarrayOrTensor: + def __call__( # type: ignore + self, img: torch.Tensor, to_pad: Optional[List[Tuple[int, int]]] = None, mode: Optional[str] = None, **kwargs + ) -> torch.Tensor: """ Args: - img: data to be transformed, assuming `img` is channel-first and - padding doesn't apply to the channel dim. - mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, - ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"`` or ``"circular"``}. - One of the listed string values or a user supplied function. Defaults to `self.mode`. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html - https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + img: data to be transformed, assuming `img` is channel-first and padding doesn't apply to the channel dim. + to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...]. + default to `self.to_pad`. + mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ - if not np.asarray(self.to_pad).any(): - # all zeros, skip padding - return img - mode = convert_pad_mode(dst=img, mode=mode or self.mode).value - pad = self._pt_pad if isinstance(img, torch.Tensor) else self._np_pad - return pad(img, self.to_pad, mode, **self.kwargs) # type: ignore - - -class SpatialPad(Transform): + to_pad_ = self.to_pad if to_pad is None else to_pad + if to_pad_ is None: + to_pad_ = self.compute_pad_width(img.shape[1:]) + mode_ = self.mode if mode is None else mode + kwargs_ = dict(self.kwargs) + kwargs_.update(kwargs) + + img_t = convert_to_tensor(data=img, track_meta=get_track_meta()) + _orig_size = img_t.shape[1:] + + # all zeros, skip padding + if np.asarray(to_pad_).any(): + if mode in ["linear_ramp", "maximum", "mean", "median", "minimum", "symmetric", "empty"]: + out = self._np_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_) + else: + try: + mode_ = convert_pad_mode(dst=img_t, mode=mode_).value + out = self._pt_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_) + # but if mode or args don't exist in pytorch, use numpy instead + except (ValueError, TypeError) as err: + if "Unsupported option" in str(err) or "unexpected keyword" in str(err): + out = self._np_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_) + else: + out = img_t + if get_track_meta(): + self.update_meta(tensor=out, to_pad=to_pad_) # type: ignore + self.push_transform(out, orig_size=_orig_size, extra_info={"padded": to_pad_}) + return out + + def update_meta(self, tensor: MetaTensor, to_pad: List[Tuple[int, int]]): + spatial_rank = max(len(tensor.affine) - 1, 1) + to_shift = [-s[0] for s in to_pad[1:]] # skipping the channel pad + mat = create_translate(spatial_rank, to_shift) + tensor.meta["affine"] = tensor.affine @ convert_to_dst_type(mat, tensor.affine)[0] + + def inverse(self, data: MetaTensor) -> MetaTensor: + transform = self.pop_transform(data) + padded = transform[TraceKeys.EXTRA_INFO]["padded"] + if padded[0][0] > 0 or padded[0][1] > 0: # slicing the channel dimension + s = padded[0][0] + e = min(max(padded[0][1], s + 1), len(data)) + data = data[s : len(data) - e] # type: ignore + roi_start = [i[0] for i in padded[1:]] + roi_end = [i - j[1] for i, j in zip(data.shape[1:], padded[1:])] + cropper = SpatialCrop(roi_start=roi_start, roi_end=roi_end) + with cropper.trace_transform(False): + return cropper(data) # type: ignore + + +class SpatialPad(Pad): """ Performs padding to the data, symmetric for all sides or all on one side for each dimension. - If input is `torch.Tensor` and mode is `constant`, `torch.nn.functional.pad` will be used. - Otherwise, `np.pad` will be used (input converted to `np.ndarray` if necessary). - - Uses np.pad so in practice, a mode needs to be provided. See numpy.lib.arraypad.pad - for additional details. - Args: spatial_size: the spatial size of output data after padding, if a dimension of the input - data size is bigger than the pad size, will not pad that dimension. + data size is larger than the pad size, will not pad that dimension. If its components have non-positive values, the corresponding size of input image will be used (no padding). for example: if the spatial size of input data is [30, 30, 30] and `spatial_size=[32, 25, -1]`, the spatial size of output data will be [32, 30, 30]. @@ -160,56 +223,37 @@ class SpatialPad(Transform): """ - backend = Pad.backend - def __init__( self, spatial_size: Union[Sequence[int], int], - method: Union[Method, str] = Method.SYMMETRIC, - mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT, + method: str = Method.SYMMETRIC, + mode: str = PytorchPadMode.CONSTANT, **kwargs, ) -> None: self.spatial_size = spatial_size self.method: Method = look_up_option(method, Method) - self.mode = mode - self.kwargs = kwargs - - def _determine_data_pad_width(self, data_shape: Sequence[int]) -> List[Tuple[int, int]]: - spatial_size = fall_back_tuple(self.spatial_size, data_shape) - if self.method == Method.SYMMETRIC: - pad_width = [] - for i, sp_i in enumerate(spatial_size): - width = max(sp_i - data_shape[i], 0) - pad_width.append((width // 2, width - (width // 2))) - return pad_width - return [(0, max(sp_i - data_shape[i], 0)) for i, sp_i in enumerate(spatial_size)] + super().__init__(mode=mode, **kwargs) - def __call__( - self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None - ) -> NdarrayOrTensor: + def compute_pad_width(self, spatial_shape: Sequence[int]) -> List[Tuple[int, int]]: """ + dynamically compute the pad width according to the spatial shape. + Args: - img: data to be transformed, assuming `img` is channel-first and - padding doesn't apply to the channel dim. - mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, - ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. - One of the listed string values or a user supplied function. Defaults to `self.mode`. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html - https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + spatial_shape: spatial shape of the original image. """ - data_pad_width = self._determine_data_pad_width(img.shape[1:]) - all_pad_width = [(0, 0)] + data_pad_width - if not np.asarray(all_pad_width).any(): - # all zeros, skip padding - return img - - padder = Pad(to_pad=all_pad_width, mode=mode or self.mode, **self.kwargs) - return padder(img) + spatial_size = fall_back_tuple(self.spatial_size, spatial_shape) + if self.method == Method.SYMMETRIC: + pad_width = [] + for i, sp_i in enumerate(spatial_size): + width = max(sp_i - spatial_shape[i], 0) + pad_width.append((width // 2, width - (width // 2))) + else: + pad_width = [(0, max(sp_i - spatial_shape[i], 0)) for i, sp_i in enumerate(spatial_size)] + return [(0, 0)] + pad_width -class BorderPad(Transform): +class BorderPad(Pad): """ Pad the input data by adding specified borders to every dimension. @@ -235,39 +279,13 @@ class BorderPad(Transform): """ - backend = Pad.backend - def __init__( - self, - spatial_border: Union[Sequence[int], int], - mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT, - **kwargs, + self, spatial_border: Union[Sequence[int], int], mode: str = PytorchPadMode.CONSTANT, **kwargs ) -> None: self.spatial_border = spatial_border - self.mode = mode - self.kwargs = kwargs + super().__init__(mode=mode, **kwargs) - def __call__( - self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None - ) -> NdarrayOrTensor: - """ - Args: - img: data to be transformed, assuming `img` is channel-first and - padding doesn't apply to the channel dim. - mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, - ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. - One of the listed string values or a user supplied function. Defaults to `self.mode`. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html - https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html - - Raises: - ValueError: When ``self.spatial_border`` does not contain ints. - ValueError: When ``self.spatial_border`` length is not one of - [1, len(spatial_shape), 2*len(spatial_shape)]. - - """ - spatial_shape = img.shape[1:] + def compute_pad_width(self, spatial_shape: Sequence[int]) -> List[Tuple[int, int]]: spatial_border = ensure_tuple(self.spatial_border) if not all(isinstance(b, int) for b in spatial_border): raise ValueError(f"self.spatial_border must contain only ints, got {spatial_border}.") @@ -284,13 +302,10 @@ def __call__( f"Unsupported spatial_border length: {len(spatial_border)}, available options are " f"[1, len(spatial_shape)={len(spatial_shape)}, 2*len(spatial_shape)={2*len(spatial_shape)}]." ) + return [(0, 0)] + data_pad_width - all_pad_width = [(0, 0)] + data_pad_width - padder = Pad(all_pad_width, mode or self.mode, **self.kwargs) - return padder(img) - -class DivisiblePad(Transform): +class DivisiblePad(Pad): """ Pad the input data, so that the spatial sizes are divisible by `k`. """ @@ -300,8 +315,8 @@ class DivisiblePad(Transform): def __init__( self, k: Union[Sequence[int], int], - mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT, - method: Union[Method, str] = Method.SYMMETRIC, + mode: str = PytorchPadMode.CONSTANT, + method: str = Method.SYMMETRIC, **kwargs, ) -> None: """ @@ -323,35 +338,120 @@ def __init__( See also :py:class:`monai.transforms.SpatialPad` """ self.k = k - self.mode: NumpyPadMode = NumpyPadMode(mode) self.method: Method = Method(method) - self.kwargs = kwargs + super().__init__(mode=mode, **kwargs) - def __call__( - self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None - ) -> NdarrayOrTensor: + def compute_pad_width(self, spatial_shape: Sequence[int]) -> List[Tuple[int, int]]: + new_size = compute_divisible_spatial_size(spatial_shape=spatial_shape, k=self.k) + spatial_pad = SpatialPad(spatial_size=new_size, method=self.method) + return spatial_pad.compute_pad_width(spatial_shape) + + +class Crop(InvertibleTransform): + """ + Perform crop operation on the input image. + + """ + + backend = [TransformBackends.TORCH] + + @staticmethod + def compute_slices( + roi_center: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_size: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_start: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_end: Union[Sequence[int], NdarrayOrTensor, None] = None, + roi_slices: Optional[Sequence[slice]] = None, + ): """ + Compute the crop slices based on specified `center & size` or `start & end`. + Args: - img: data to be transformed, assuming `img` is channel-first - and padding doesn't apply to the channel dim. - mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, - ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. - One of the listed string values or a user supplied function. Defaults to `self.mode`. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html - https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + roi_center: voxel coordinates for center of the crop ROI. + roi_size: size of the crop ROI, if a dimension of ROI size is larger than image size, + will not crop that dimension of the image. + roi_start: voxel coordinates for start of the crop ROI. + roi_end: voxel coordinates for end of the crop ROI, if a coordinate is out of image, + use the end coordinate of image. + roi_slices: list of slices for each of the spatial dimensions. """ - new_size = compute_divisible_spatial_size(spatial_shape=img.shape[1:], k=self.k) - spatial_pad = SpatialPad(spatial_size=new_size, method=self.method, mode=mode or self.mode, **self.kwargs) + roi_start_t: torch.Tensor - return spatial_pad(img) + if roi_slices: + if not all(s.step is None or s.step == 1 for s in roi_slices): + raise ValueError("only slice steps of 1/None are currently supported") + return list(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) + roi_size_t = convert_to_tensor(data=roi_size, dtype=torch.int16, wrap_sequence=True) + _zeros = torch.zeros_like(roi_center_t) + half = ( + torch.divide(roi_size_t, 2, rounding_mode="floor") + if pytorch_after(1, 8) + else torch.floor_divide(roi_size_t, 2) + ) + roi_start_t = torch.maximum(roi_center_t - half, _zeros) + roi_end_t = torch.maximum(roi_start_t + roi_size_t, roi_start_t) + 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 [slice(int(roi_start_t.item()), int(roi_end_t.item()))] + else: + return [slice(int(s), int(e)) for s, e in zip(roi_start_t.tolist(), roi_end_t.tolist())] + def __call__(self, img: torch.Tensor, slices: Tuple[slice, ...]) -> torch.Tensor: # type: ignore + """ + Apply the transform to `img`, assuming `img` is channel-first and + slicing doesn't apply to the channel dim. -class SpatialCrop(Transform): + """ + orig_size = img.shape[1:] + slices_ = list(slices) + sd = len(img.shape[1:]) # spatial dims + if len(slices_) < sd: + slices_ += [slice(None)] * (sd - len(slices_)) + # Add in the channel (no cropping) + slices = tuple([slice(None)] + slices_[:sd]) + + img_t: MetaTensor = convert_to_tensor(data=img, track_meta=get_track_meta()) + _orig_size = img_t.shape[1:] + img_t = img_t[slices] # type: ignore + if get_track_meta(): + self.update_meta(tensor=img_t, slices=slices) + cropped_from_start = np.asarray([s.indices(o)[0] for s, o in zip(slices[1:], orig_size)]) + cropped_from_end = np.asarray(orig_size) - img_t.shape[1:] - cropped_from_start + cropped = list(chain(*zip(cropped_from_start.tolist(), cropped_from_end.tolist()))) + self.push_transform(img_t, orig_size=_orig_size, extra_info={"cropped": cropped}) + return img_t + + def update_meta(self, tensor: MetaTensor, slices: Tuple[slice, ...]): + spatial_rank = max(len(tensor.affine) - 1, 1) + to_shift = [s.start if s.start is not None else 0 for s in ensure_tuple(slices)[1:]] + mat = create_translate(spatial_rank, to_shift) + tensor.meta["affine"] = tensor.affine @ convert_to_dst_type(mat, tensor.affine)[0] + + def inverse(self, img: MetaTensor) -> MetaTensor: + transform = self.pop_transform(img) + cropped = transform[TraceKeys.EXTRA_INFO]["cropped"] + # the amount we pad is equal to the amount we cropped in each direction + inverse_transform = BorderPad(cropped) + # Apply inverse transform + with inverse_transform.trace_transform(False): + return inverse_transform(img) # type: ignore + + +class SpatialCrop(Crop): """ General purpose cropper to produce sub-volume region of interest (ROI). - If a dimension of the expected ROI size is bigger than the input image size, will not crop that dimension. + If a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. It can support to crop ND spatial (channel-first) data. @@ -362,8 +462,6 @@ class SpatialCrop(Transform): - the start and end coordinates of the ROI """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] - def __init__( self, roi_center: Union[Sequence[int], NdarrayOrTensor, None] = None, @@ -375,85 +473,59 @@ def __init__( """ Args: roi_center: voxel coordinates for center of the crop ROI. - roi_size: size of the crop ROI, if a dimension of ROI size is bigger than image size, + roi_size: size of the crop ROI, if a dimension of ROI size is larger than image size, will not crop that dimension of the image. roi_start: voxel coordinates for start of the crop ROI. roi_end: voxel coordinates for end of the crop ROI, if a coordinate is out of image, use the end coordinate of image. roi_slices: list of slices for each of the spatial dimensions. """ - roi_start_torch: torch.Tensor - - if roi_slices: - if not all(s.step is None or s.step == 1 for s in roi_slices): - raise ValueError("Only slice steps of 1/None are currently supported") - self.slices = list(roi_slices) - else: - if roi_center is not None and roi_size is not None: - roi_center, *_ = convert_data_type( - data=roi_center, output_type=torch.Tensor, dtype=torch.int16, wrap_sequence=True - ) - roi_size, *_ = convert_to_dst_type(src=roi_size, dst=roi_center, wrap_sequence=True) - _zeros = torch.zeros_like(roi_center) - roi_start_torch = maximum(roi_center - floor_divide(roi_size, 2), _zeros) # type: ignore - roi_end_torch = maximum(roi_start_torch + roi_size, roi_start_torch) - 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_torch, *_ = convert_data_type( - data=roi_start, output_type=torch.Tensor, dtype=torch.int16, wrap_sequence=True - ) - roi_start_torch = maximum(roi_start_torch, torch.zeros_like(roi_start_torch)) # type: ignore - roi_end_torch, *_ = convert_to_dst_type(src=roi_end, dst=roi_start_torch, wrap_sequence=True) - roi_end_torch = maximum(roi_end_torch, roi_start_torch) - # convert to slices (accounting for 1d) - if roi_start_torch.numel() == 1: - self.slices = [slice(int(roi_start_torch.item()), int(roi_end_torch.item()))] - else: - self.slices = [slice(int(s), int(e)) for s, e in zip(roi_start_torch.tolist(), roi_end_torch.tolist())] + self.slices = self.compute_slices( + roi_center=roi_center, roi_size=roi_size, roi_start=roi_start, roi_end=roi_end, roi_slices=roi_slices + ) - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor) -> torch.Tensor: # type: ignore """ Apply the transform to `img`, assuming `img` is channel-first and slicing doesn't apply to the channel dim. + """ - sd = min(len(self.slices), len(img.shape[1:])) # spatial dims - slices = [slice(None)] + self.slices[:sd] - return img[tuple(slices)] + return super().__call__(img=img, slices=self.slices) -class CenterSpatialCrop(Transform): +class CenterSpatialCrop(Crop): """ Crop at the center of image with specified ROI size. - If a dimension of the expected ROI size is bigger than the input image size, will not crop that dimension. + If a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. Args: roi_size: the spatial size of the crop region e.g. [224,224,128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. """ - backend = SpatialCrop.backend - def __init__(self, roi_size: Union[Sequence[int], int]) -> None: self.roi_size = roi_size - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def compute_slices(self, spatial_size: Sequence[int]): # type: ignore + roi_size = fall_back_tuple(self.roi_size, spatial_size) + roi_center = [i // 2 for i in spatial_size] + return super().compute_slices(roi_center=roi_center, roi_size=roi_size) + + def __call__(self, img: torch.Tensor) -> torch.Tensor: # type: ignore """ Apply the transform to `img`, assuming `img` is channel-first and slicing doesn't apply to the channel dim. + """ - roi_size = fall_back_tuple(self.roi_size, img.shape[1:]) - center = [i // 2 for i in img.shape[1:]] - cropper = SpatialCrop(roi_center=center, roi_size=roi_size) - return cropper(img) + return super().__call__(img=img, slices=self.compute_slices(img.shape[1:])) -class CenterScaleCrop(Transform): +class CenterScaleCrop(Crop): """ Crop at the center of image with specified scale of ROI size. @@ -463,32 +535,30 @@ class CenterScaleCrop(Transform): """ - backend = CenterSpatialCrop.backend - def __init__(self, roi_scale: Union[Sequence[float], float]): self.roi_scale = roi_scale - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor) -> torch.Tensor: # type: ignore img_size = img.shape[1:] ndim = len(img_size) roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.roi_scale, ndim), img_size)] - sp_crop = CenterSpatialCrop(roi_size=roi_size) - return sp_crop(img=img) + cropper = CenterSpatialCrop(roi_size=roi_size) + return super().__call__(img=img, slices=cropper.compute_slices(img.shape[1:])) -class RandSpatialCrop(Randomizable, Transform): +class RandSpatialCrop(Randomizable, Crop): """ Crop image with random size or specific size ROI. It can crop at a random position as center or at the image center. And allows to set the minimum and maximum size to limit the randomly generated ROI. - Note: even `random_size=False`, if a dimension of the expected ROI size is bigger than the input image size, + Note: even `random_size=False`, if a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. Args: roi_size: if `random_size` is True, it specifies the minimum crop region. if `random_size` is False, it specifies the expected ROI size to crop. e.g. [224, 224, 128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -500,8 +570,6 @@ class RandSpatialCrop(Randomizable, Transform): if True, the actual size is sampled from `randint(roi_size, max_roi_size + 1)`. """ - backend = CenterSpatialCrop.backend - def __init__( self, roi_size: Union[Sequence[int], int], @@ -514,31 +582,33 @@ def __init__( self.random_center = random_center self.random_size = random_size self._size: Optional[Sequence[int]] = None - self._slices: Optional[Tuple[slice, ...]] = None + self._slices: Tuple[slice, ...] def randomize(self, img_size: Sequence[int]) -> None: self._size = fall_back_tuple(self.roi_size, img_size) if self.random_size: max_size = img_size if self.max_roi_size is None else fall_back_tuple(self.max_roi_size, img_size) if any(i > j for i, j in zip(self._size, max_size)): - raise ValueError(f"min ROI size: {self._size} is bigger than max ROI size: {max_size}.") + raise ValueError(f"min ROI size: {self._size} is larger than max ROI size: {max_size}.") self._size = tuple(self.R.randint(low=self._size[i], high=max_size[i] + 1) for i in range(len(img_size))) if self.random_center: valid_size = get_valid_patch_size(img_size, self._size) - self._slices = (slice(None),) + get_random_patch(img_size, valid_size, self.R) + self._slices = get_random_patch(img_size, valid_size, self.R) - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor: # type: ignore """ Apply the transform to `img`, assuming `img` is channel-first and slicing doesn't apply to the channel dim. + """ - self.randomize(img.shape[1:]) + if randomize: + self.randomize(img.shape[1:]) if self._size is None: raise RuntimeError("self._size not specified.") if self.random_center: - return img[self._slices] + return super().__call__(img=img, slices=self._slices) cropper = CenterSpatialCrop(self._size) - return cropper(img) + return super().__call__(img=img, slices=cropper.compute_slices(img.shape[1:])) class RandScaleCrop(RandSpatialCrop): @@ -573,36 +643,43 @@ def __init__( self.roi_scale = roi_scale self.max_roi_scale = max_roi_scale - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: - """ - Apply the transform to `img`, assuming `img` is channel-first and - slicing doesn't apply to the channel dim. - """ - img_size = img.shape[1:] + def get_max_roi_size(self, img_size): ndim = len(img_size) self.roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.roi_scale, ndim), img_size)] if self.max_roi_scale is not None: self.max_roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.max_roi_scale, ndim), img_size)] else: self.max_roi_size = None - return super().__call__(img=img) + def randomize(self, img_size: Sequence[int]) -> None: + self.get_max_roi_size(img_size) + super().randomize(img_size) + + def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor: # type: ignore + """ + Apply the transform to `img`, assuming `img` is channel-first and + slicing doesn't apply to the channel dim. -class RandSpatialCropSamples(Randomizable, Transform): + """ + self.get_max_roi_size(img.shape[1:]) + return super().__call__(img=img, randomize=randomize) + + +class RandSpatialCropSamples(Randomizable, TraceableTransform): """ Crop image with random size or specific size ROI to generate a list of N samples. It can crop at a random position as center or at the image center. And allows to set the minimum size to limit the randomly generated ROI. It will return a list of cropped images. - Note: even `random_size=False`, if a dimension of the expected ROI size is bigger than the input image size, + Note: even `random_size=False`, if a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. Args: roi_size: if `random_size` is True, it specifies the minimum crop region. if `random_size` is False, it specifies the expected ROI size to crop. e.g. [224, 224, 128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -644,15 +721,23 @@ def set_random_state( def randomize(self, data: Optional[Any] = None) -> None: pass - def __call__(self, img: NdarrayOrTensor) -> List[NdarrayOrTensor]: + def __call__(self, img: torch.Tensor) -> List[torch.Tensor]: """ Apply the transform to `img`, assuming `img` is channel-first and cropping doesn't change the channel dim. """ - return [self.cropper(img) for _ in range(self.num_samples)] - - -class CropForeground(Transform): + ret = [] + orig_size = img.shape[1:] + for i in range(self.num_samples): + cropped = self.cropper(img) + if get_track_meta(): + cropped.meta[Key.PATCH_INDEX] = i # type: ignore + self.push_transform(cropped, orig_size=orig_size, extra_info=self.pop_transform(cropped, check=False)) + ret.append(cropped) + return ret + + +class CropForeground(Crop): """ Crop an image using a bounding box. The bounding box is generated by selecting foreground using select_fn at channels channel_indices. margin is added in each spatial dimension of the bounding box. @@ -684,8 +769,6 @@ def threshold_at_one(x): """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] - def __init__( self, select_fn: Callable = is_positive, @@ -694,8 +777,8 @@ def __init__( allow_smaller: bool = True, return_coords: bool = False, k_divisible: Union[Sequence[int], int] = 1, - mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = NumpyPadMode.CONSTANT, - **np_kwargs, + mode: str = PytorchPadMode.CONSTANT, + **pad_kwargs, ) -> None: """ Args: @@ -704,7 +787,7 @@ def __init__( of image. if None, select foreground on the whole image. margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims. allow_smaller: when computing box size with `margin`, whether allow the image size to be smaller - than box size, default to `True`. if the margined size is bigger than image size, will pad with + than box size, default to `True`. if the margined size is larger than image size, will pad with specified `mode`. return_coords: whether return the coordinates of spatial bounding box for foreground. k_divisible: make each spatial dimension to be divisible by k, default to 1. @@ -715,8 +798,8 @@ def __init__( One of the listed string values or a user supplied function. Defaults to ``"constant"``. See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ self.select_fn = select_fn @@ -725,10 +808,9 @@ def __init__( self.allow_smaller = allow_smaller self.return_coords = return_coords self.k_divisible = k_divisible - self.mode: NumpyPadMode = look_up_option(mode, NumpyPadMode) - self.np_kwargs = np_kwargs + self.padder = Pad(mode=mode, **pad_kwargs) - def compute_bounding_box(self, img: NdarrayOrTensor): + def compute_bounding_box(self, img: torch.Tensor): """ Compute the start points and end points of bounding box to crop. And adjust bounding box coords to be divisible by `k`. @@ -748,36 +830,51 @@ def compute_bounding_box(self, img: NdarrayOrTensor): return box_start_, box_end_ def crop_pad( - self, - img: NdarrayOrTensor, - box_start: np.ndarray, - box_end: np.ndarray, - mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None, + self, img: torch.Tensor, box_start: np.ndarray, box_end: np.ndarray, mode: Optional[str] = None, **pad_kwargs ): """ Crop and pad based on the bounding box. """ - cropped = SpatialCrop(roi_start=box_start, roi_end=box_end)(img) + slices = self.compute_slices(roi_start=box_start, roi_end=box_end) + cropped = super().__call__(img=img, slices=slices) pad_to_start = np.maximum(-box_start, 0) pad_to_end = np.maximum(box_end - np.asarray(img.shape[1:]), 0) pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - return BorderPad(spatial_border=pad, mode=mode or self.mode, **self.np_kwargs)(cropped) - - def __call__(self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, str]] = None): + pad_width = BorderPad(spatial_border=pad).compute_pad_width(cropped.shape[1:]) + ret = self.padder.__call__(img=cropped, to_pad=pad_width, mode=mode, **pad_kwargs) + # combine the traced cropping and padding into one transformation + # by taking the padded info and placing it in a key inside the crop info. + if get_track_meta(): + ret_: MetaTensor = ret # type: ignore + app_op = ret_.applied_operations.pop(-1) + ret_.applied_operations[-1][TraceKeys.EXTRA_INFO]["pad_info"] = app_op + return ret + + def __call__(self, img: torch.Tensor, mode: Optional[str] = None, **pad_kwargs): # type: ignore """ Apply the transform to `img`, assuming `img` is channel-first and slicing doesn't change the channel dim. """ box_start, box_end = self.compute_bounding_box(img) - cropped = self.crop_pad(img, box_start, box_end, mode) + cropped = self.crop_pad(img, box_start, box_end, mode, **pad_kwargs) if self.return_coords: return cropped, box_start, box_end return cropped + def inverse(self, img: MetaTensor) -> MetaTensor: + transform = self.get_most_recent_transform(img) + # we moved the padding info in the forward, so put it back for the inverse + pad_info = transform[TraceKeys.EXTRA_INFO].pop("pad_info") + img.applied_operations.append(pad_info) + # first inverse the padder + inv = self.padder.inverse(img) + # and then inverse the cropper (self) + return super().inverse(inv) + -class RandWeightedCrop(Randomizable, Transform): +class RandWeightedCrop(Randomizable, TraceableTransform): """ Samples a list of `num_samples` image patches according to the provided `weight_map`. @@ -808,13 +905,16 @@ def randomize(self, weight_map: NdarrayOrTensor) -> None: spatial_size=self.spatial_size, w=weight_map[0], n_samples=self.num_samples, r_state=self.R ) # using only the first channel as weight map - def __call__(self, img: NdarrayOrTensor, weight_map: Optional[NdarrayOrTensor] = None) -> List[NdarrayOrTensor]: + def __call__( + self, img: torch.Tensor, weight_map: Optional[NdarrayOrTensor] = None, randomize: bool = True + ) -> List[torch.Tensor]: """ Args: img: input image to sample patches from. assuming `img` is a channel-first array. 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, ...)` + randomize: whether to execute random operations, default to `True`. Returns: A list of image patches @@ -826,16 +926,23 @@ def __call__(self, img: NdarrayOrTensor, weight_map: Optional[NdarrayOrTensor] = if img.shape[1:] != weight_map.shape[1:]: raise ValueError(f"image and weight map spatial shape mismatch: {img.shape[1:]} vs {weight_map.shape[1:]}.") - self.randomize(weight_map) + if randomize: + self.randomize(weight_map) _spatial_size = fall_back_tuple(self.spatial_size, weight_map.shape[1:]) - results: List[NdarrayOrTensor] = [] - for center in self.centers: - cropper = SpatialCrop(roi_center=center, roi_size=_spatial_size) - results.append(cropper(img)) + results: List[torch.Tensor] = [] + orig_size = img.shape[1:] + for i, center in enumerate(self.centers): + cropped = SpatialCrop(roi_center=center, roi_size=_spatial_size)(img) + if get_track_meta(): + ret_: MetaTensor = cropped # type: ignore + ret_.meta[Key.PATCH_INDEX] = i + ret_.meta["crop_center"] = center + self.push_transform(ret_, orig_size=orig_size, extra_info=self.pop_transform(ret_, check=False)) + results.append(cropped) return results -class RandCropByPosNegLabel(Randomizable, Transform): +class RandCropByPosNegLabel(Randomizable, TraceableTransform): """ Crop random fixed sized regions with the center being a foreground or background voxel based on the Pos Neg Ratio. @@ -848,13 +955,15 @@ class RandCropByPosNegLabel(Randomizable, Transform): [0, 0, 0, 0, 0], [0, 0, 0]] [0, 0, 0]] [0, 0, 0, 0, 0]]] - If a dimension of the expected spatial size is bigger than the input image size, + If a dimension of the expected spatial size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than expected size, and the cropped results of several images may not have exactly same shape. + And if the crop ROI is partly out of the image, will automatically adjust the crop center to ensure the + valid crop ROI. Args: spatial_size: the spatial size of the crop region e.g. [224, 224, 128]. - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. if its components have non-positive values, the corresponding size of `label` will be used. for example: if the spatial size of input data is [40, 40, 40] and `spatial_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -888,22 +997,22 @@ class RandCropByPosNegLabel(Randomizable, Transform): """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + backend = SpatialCrop.backend def __init__( self, spatial_size: Union[Sequence[int], int], - label: Optional[NdarrayOrTensor] = None, + label: Optional[torch.Tensor] = None, pos: float = 1.0, neg: float = 1.0, num_samples: int = 1, - image: Optional[NdarrayOrTensor] = None, + image: Optional[torch.Tensor] = None, image_threshold: float = 0.0, fg_indices: Optional[NdarrayOrTensor] = None, bg_indices: Optional[NdarrayOrTensor] = None, allow_smaller: bool = False, ) -> None: - self.spatial_size = ensure_tuple(spatial_size) + self.spatial_size = spatial_size self.label = label if pos < 0 or neg < 0: raise ValueError(f"pos and neg must be nonnegative, got pos={pos} neg={neg}.") @@ -920,12 +1029,11 @@ def __init__( def randomize( self, - label: NdarrayOrTensor, + label: torch.Tensor, fg_indices: Optional[NdarrayOrTensor] = None, bg_indices: Optional[NdarrayOrTensor] = None, - image: Optional[NdarrayOrTensor] = None, + image: Optional[torch.Tensor] = None, ) -> None: - self.spatial_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) if fg_indices is None or bg_indices is None: if self.fg_indices is not None and self.bg_indices is not None: fg_indices_ = self.fg_indices @@ -948,12 +1056,13 @@ def randomize( def __call__( self, - img: NdarrayOrTensor, - label: Optional[NdarrayOrTensor] = None, - image: Optional[NdarrayOrTensor] = None, + img: torch.Tensor, + label: Optional[torch.Tensor] = None, + image: Optional[torch.Tensor] = None, fg_indices: Optional[NdarrayOrTensor] = None, bg_indices: Optional[NdarrayOrTensor] = None, - ) -> List[NdarrayOrTensor]: + randomize: bool = True, + ) -> List[torch.Tensor]: """ Args: img: input data to crop samples from based on the pos/neg ratio of `label` and `image`. @@ -966,6 +1075,7 @@ def __call__( need to provide `fg_indices` and `bg_indices` together. bg_indices: background indices to randomly select crop centers, need to provide `fg_indices` and `bg_indices` together. + randomize: whether to execute the random operations, default to `True`. """ if label is None: @@ -975,17 +1085,24 @@ def __call__( if image is None: image = self.image - self.randomize(label, fg_indices, bg_indices, image) - results: List[NdarrayOrTensor] = [] + if randomize: + self.randomize(label, fg_indices, bg_indices, image) + results: List[torch.Tensor] = [] + orig_size = img.shape[1:] if self.centers is not None: - for center in self.centers: - cropper = SpatialCrop(roi_center=center, roi_size=self.spatial_size) - results.append(cropper(img)) - + for i, center in enumerate(self.centers): + roi_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) + cropped = SpatialCrop(roi_center=center, roi_size=roi_size)(img) + if get_track_meta(): + ret_: MetaTensor = cropped # type: ignore + ret_.meta[Key.PATCH_INDEX] = i + ret_.meta["crop_center"] = center + self.push_transform(ret_, orig_size=orig_size, extra_info=self.pop_transform(ret_, check=False)) + results.append(cropped) return results -class RandCropByLabelClasses(Randomizable, Transform): +class RandCropByLabelClasses(Randomizable, TraceableTransform): """ Crop random fixed sized regions with the center being a class based on the specified ratios of every class. The label data can be One-Hot format array or Argmax data. And will return a list of arrays for all the @@ -1018,13 +1135,15 @@ class RandCropByLabelClasses(Randomizable, Transform): [0, 1, 3], [1, 2, 1], [0, 0, 0]] [1, 3, 0]] - If a dimension of the expected spatial size is bigger than the input image size, + If a dimension of the expected spatial size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than expected size, and the cropped results of several images may not have exactly same shape. + And if the crop ROI is partly out of the image, will automatically adjust the crop center to ensure the + valid crop ROI. Args: spatial_size: the spatial size of the crop region e.g. [224, 224, 128]. - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. if its components have non-positive values, the corresponding size of `label` will be used. for example: if the spatial size of input data is [40, 40, 40] and `spatial_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -1047,21 +1166,21 @@ class RandCropByLabelClasses(Randomizable, Transform): """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + backend = SpatialCrop.backend def __init__( self, spatial_size: Union[Sequence[int], int], ratios: Optional[List[Union[float, int]]] = None, - label: Optional[NdarrayOrTensor] = None, + label: Optional[torch.Tensor] = None, num_classes: Optional[int] = None, num_samples: int = 1, - image: Optional[NdarrayOrTensor] = None, + image: Optional[torch.Tensor] = None, image_threshold: float = 0.0, indices: Optional[List[NdarrayOrTensor]] = None, allow_smaller: bool = False, ) -> None: - self.spatial_size = ensure_tuple(spatial_size) + self.spatial_size = spatial_size self.ratios = ratios self.label = label self.num_classes = num_classes @@ -1073,12 +1192,8 @@ def __init__( self.allow_smaller = allow_smaller def randomize( - self, - label: NdarrayOrTensor, - indices: Optional[List[NdarrayOrTensor]] = None, - image: Optional[NdarrayOrTensor] = None, + self, label: torch.Tensor, indices: Optional[List[NdarrayOrTensor]] = None, image: Optional[torch.Tensor] = None ) -> None: - self.spatial_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) indices_: Sequence[NdarrayOrTensor] if indices is None: if self.indices is not None: @@ -1093,11 +1208,12 @@ def randomize( def __call__( self, - img: NdarrayOrTensor, - label: Optional[NdarrayOrTensor] = None, - image: Optional[NdarrayOrTensor] = None, + img: torch.Tensor, + label: Optional[torch.Tensor] = None, + image: Optional[torch.Tensor] = None, indices: Optional[List[NdarrayOrTensor]] = None, - ) -> List[NdarrayOrTensor]: + randomize: bool = True, + ) -> List[torch.Tensor]: """ Args: img: input data to crop samples from based on the ratios of every class, assumes `img` is a @@ -1106,6 +1222,7 @@ def __call__( image: optional image data to help select valid area, can be same as `img` or another image array. use ``image > image_threshold`` to select the centers only in valid region. if None, use `self.image`. indices: list of indices for every class in the image, used to randomly select crop centers. + randomize: whether to execute the random operations, default to `True`. """ if label is None: @@ -1115,17 +1232,25 @@ def __call__( if image is None: image = self.image - self.randomize(label, indices, image) - results: List[NdarrayOrTensor] = [] + if randomize: + self.randomize(label, indices, image) + results: List[torch.Tensor] = [] + orig_size = img.shape[1:] if self.centers is not None: - for center in self.centers: - cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) - results.append(cropper(img)) + for i, center in enumerate(self.centers): + roi_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) + cropped = SpatialCrop(roi_center=tuple(center), roi_size=roi_size)(img) + if get_track_meta(): + ret_: MetaTensor = cropped # type: ignore + ret_.meta[Key.PATCH_INDEX] = i + ret_.meta["crop_center"] = center + self.push_transform(ret_, orig_size=orig_size, extra_info=self.pop_transform(ret_, check=False)) + results.append(cropped) return results -class ResizeWithPadOrCrop(Transform): +class ResizeWithPadOrCrop(InvertibleTransform): """ Resize an image to a target spatial size by either centrally cropping the image or padding it evenly with a user-specified mode. @@ -1135,14 +1260,16 @@ class ResizeWithPadOrCrop(Transform): Args: spatial_size: the spatial size of output data after padding or crop. If has non-positive values, the corresponding size of input image will be used (no padding). - mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, - ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - One of the listed string values or a user supplied function for padding. Defaults to ``"constant"``. - See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html method: {``"symmetric"``, ``"end"``} Pad image symmetrically on every side or only pad at the end sides. Defaults to ``"symmetric"``. - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ @@ -1151,25 +1278,52 @@ class ResizeWithPadOrCrop(Transform): def __init__( self, spatial_size: Union[Sequence[int], int], - mode: Union[NumpyPadMode, str] = NumpyPadMode.CONSTANT, - method: Union[Method, str] = Method.SYMMETRIC, - **np_kwargs, + method: str = Method.SYMMETRIC, + mode: str = PytorchPadMode.CONSTANT, + **pad_kwargs, ): - self.padder = SpatialPad(spatial_size=spatial_size, method=method, mode=mode, **np_kwargs) + self.padder = SpatialPad(spatial_size=spatial_size, method=method, mode=mode, **pad_kwargs) self.cropper = CenterSpatialCrop(roi_size=spatial_size) - def __call__(self, img: NdarrayOrTensor, mode: Optional[Union[NumpyPadMode, str]] = None) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, mode: Optional[str] = None, **pad_kwargs) -> torch.Tensor: # type: ignore """ Args: img: data to pad or crop, assuming `img` is channel-first and padding or cropping doesn't apply to the channel dim. - mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, - ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - One of the listed string values or a user supplied function for padding. - If None, defaults to the ``mode`` in construction. + mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. + """ - return self.padder(self.cropper(img), mode=mode) + orig_size = img.shape[1:] + ret = self.padder(self.cropper(img), mode=mode, **pad_kwargs) + # remove the individual info and combine + if get_track_meta(): + ret_: MetaTensor = ret # type: ignore + pad_info = ret_.applied_operations.pop(-1) + crop_info = ret_.applied_operations.pop(-1) + self.push_transform(ret_, orig_size=orig_size, extra_info={"pad_info": pad_info, "crop_info": crop_info}) + return ret + + def inverse(self, img: MetaTensor) -> MetaTensor: + transform = self.pop_transform(img) + return self.inverse_transform(img, transform) + + def inverse_transform(self, img: MetaTensor, transform) -> MetaTensor: + # we joined the cropping and padding, so put them back before calling the inverse + crop_info = transform[TraceKeys.EXTRA_INFO].pop("crop_info") + pad_info = transform[TraceKeys.EXTRA_INFO].pop("pad_info") + img.applied_operations.append(crop_info) + img.applied_operations.append(pad_info) + # first inverse the padder + inv = self.padder.inverse(img) + # and then inverse the cropper (self) + return self.cropper.inverse(inv) class BoundingRect(Transform): diff --git a/monai/transforms/croppad/batch.py b/monai/transforms/croppad/batch.py index 52d0c7be3b..34945b6b84 100644 --- a/monai/transforms/croppad/batch.py +++ b/monai/transforms/croppad/batch.py @@ -14,15 +14,16 @@ """ from copy import deepcopy -from typing import Any, Dict, Hashable, Union +from typing import Any, Dict, Hashable import numpy as np import torch +from monai.data.meta_tensor import MetaTensor from monai.data.utils import list_data_collate from monai.transforms.croppad.array import CenterSpatialCrop, SpatialPad from monai.transforms.inverse import InvertibleTransform -from monai.utils.enums import Method, NumpyPadMode, TraceKeys +from monai.utils.enums import Method, PytorchPadMode, TraceKeys __all__ = ["PadListDataCollate"] @@ -57,20 +58,15 @@ class PadListDataCollate(InvertibleTransform): Args: method: padding method (see :py:class:`monai.transforms.SpatialPad`) mode: padding mode (see :py:class:`monai.transforms.SpatialPad`) - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ - def __init__( - self, - method: Union[Method, str] = Method.SYMMETRIC, - mode: Union[NumpyPadMode, str] = NumpyPadMode.CONSTANT, - **np_kwargs, - ) -> None: + def __init__(self, method: str = Method.SYMMETRIC, mode: str = PytorchPadMode.CONSTANT, **kwargs) -> None: self.method = method self.mode = mode - self.np_kwargs = np_kwargs + self.kwargs = kwargs def __call__(self, batch: Any): """ @@ -80,7 +76,8 @@ def __call__(self, batch: Any): # data is either list of dicts or list of lists is_list_of_dicts = isinstance(batch[0], dict) # loop over items inside of each element in a batch - for key_or_idx in batch[0].keys() if is_list_of_dicts else range(len(batch[0])): + batch_item = tuple(batch[0].keys()) if is_list_of_dicts else range(len(batch[0])) + for key_or_idx in batch_item: # calculate max size of each dimension max_shapes = [] for elem in batch: @@ -96,7 +93,7 @@ def __call__(self, batch: Any): continue # Use `SpatialPad` to match sizes, Default params are central padding, padding with 0's - padder = SpatialPad(spatial_size=max_shape, method=self.method, mode=self.mode, **self.np_kwargs) + padder = SpatialPad(spatial_size=max_shape, method=self.method, mode=self.mode, **self.kwargs) for idx, batch_i in enumerate(batch): orig_size = batch_i[key_or_idx].shape[1:] padded = padder(batch_i[key_or_idx]) @@ -116,13 +113,18 @@ def inverse(data: dict) -> Dict[Hashable, np.ndarray]: d = deepcopy(data) for key in d: - transform_key = InvertibleTransform.trace_key(key) - if transform_key in d: - transform = d[transform_key][-1] - if not isinstance(transform, Dict): - continue - if transform.get(TraceKeys.CLASS_NAME) == PadListDataCollate.__name__: - d[key] = CenterSpatialCrop(transform.get("orig_size", -1))(d[key]) # fallback to image size - # remove transform - d[transform_key].pop() + transforms = None + if isinstance(d[key], MetaTensor): + transforms = d[key].applied_operations + else: + transform_key = InvertibleTransform.trace_key(key) + if transform_key in d: + transforms = d[transform_key] + if not transforms or not isinstance(transforms[-1], Dict): + continue + if transforms[-1].get(TraceKeys.CLASS_NAME) == PadListDataCollate.__name__: + xform = transforms.pop() + cropping = CenterSpatialCrop(xform.get(TraceKeys.ORIG_SIZE, -1)) + with cropping.trace_transform(False): + d[key] = cropping(d[key]) # fallback to image size return d diff --git a/monai/transforms/croppad/dictionary.py b/monai/transforms/croppad/dictionary.py index 19ebe40b46..a3310a10d1 100644 --- a/monai/transforms/croppad/dictionary.py +++ b/monai/transforms/croppad/dictionary.py @@ -15,50 +15,47 @@ Class names are ended with 'd' to denote dictionary-based transforms. """ -import contextlib from copy import deepcopy -from enum import Enum -from itertools import chain -from math import ceil, floor -from typing import Any, Callable, Dict, Hashable, List, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, Hashable, List, Mapping, Optional, Sequence, Union import numpy as np +import torch -from monai.config import IndexSelection, KeysCollection +from monai.config import IndexSelection, KeysCollection, SequenceStr from monai.config.type_definitions import NdarrayOrTensor -from monai.data.utils import get_random_patch, get_valid_patch_size +from monai.data.meta_tensor import MetaTensor from monai.transforms.croppad.array import ( BorderPad, BoundingRect, + CenterScaleCrop, CenterSpatialCrop, + Crop, CropForeground, DivisiblePad, + Pad, RandCropByLabelClasses, RandCropByPosNegLabel, + RandScaleCrop, + RandSpatialCrop, + RandSpatialCropSamples, + RandWeightedCrop, ResizeWithPadOrCrop, SpatialCrop, SpatialPad, ) from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import MapTransform, Randomizable -from monai.transforms.utils import ( - allow_missing_keys_mode, - generate_label_classes_crop_centers, - generate_pos_neg_label_crop_centers, - is_positive, - map_binary_to_indices, - map_classes_to_indices, - weighted_patch_samples, -) -from monai.utils import ImageMetaKey as Key -from monai.utils import Method, NumpyPadMode, PytorchPadMode, ensure_tuple, ensure_tuple_rep, fall_back_tuple -from monai.utils.enums import PostFix, TraceKeys +from monai.transforms.utils import is_positive +from monai.utils import MAX_SEED, Method, PytorchPadMode, ensure_tuple_rep +from monai.utils.deprecate_utils import deprecated_arg __all__ = [ - "PadModeSequence", + "Padd", "SpatialPadd", "BorderPadd", "DivisiblePadd", + "Cropd", + "RandCropd", "SpatialCropd", "CenterSpatialCropd", "CenterScaleCropd", @@ -70,12 +67,19 @@ "RandCropByPosNegLabeld", "ResizeWithPadOrCropd", "BoundingRectd", + "RandCropByLabelClassesd", + "PadD", + "PadDict", "SpatialPadD", "SpatialPadDict", "BorderPadD", "BorderPadDict", "DivisiblePadD", "DivisiblePadDict", + "CropD", + "CropDict", + "RandCropD", + "RandCropDict", "SpatialCropD", "SpatialCropDict", "CenterSpatialCropD", @@ -98,30 +102,71 @@ "ResizeWithPadOrCropDict", "BoundingRectD", "BoundingRectDict", - "RandCropByLabelClassesd", "RandCropByLabelClassesD", "RandCropByLabelClassesDict", ] -NumpyPadModeSequence = Union[Sequence[Union[NumpyPadMode, str]], NumpyPadMode, str] -PadModeSequence = Union[Sequence[Union[NumpyPadMode, PytorchPadMode, str]], NumpyPadMode, PytorchPadMode, str] -DEFAULT_POST_FIX = PostFix.meta() +class Padd(MapTransform, InvertibleTransform): + """ + Dictionary-based wrapper of :py:class:`monai.transforms.Pad`. -class SpatialPadd(MapTransform, InvertibleTransform): + """ + + backend = Pad.backend + + def __init__( + self, + keys: KeysCollection, + padder: Pad, + mode: SequenceStr = PytorchPadMode.CONSTANT, + allow_missing_keys: bool = False, + ) -> None: + """ + Args: + keys: keys of the corresponding items to be transformed. + See also: :py:class:`monai.transforms.compose.MapTransform` + padder: pad transform for the input image. + mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html + It also can be a sequence of string, each element corresponds to a key in ``keys``. + allow_missing_keys: don't raise exception if key is missing. + + """ + super().__init__(keys, allow_missing_keys) + self.padder = padder + self.mode = ensure_tuple_rep(mode, len(self.keys)) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + for key, m in self.key_iterator(d, self.mode): + d[key] = self.padder(d[key], mode=m) + return d + + def inverse(self, data: Mapping[Hashable, MetaTensor]) -> Dict[Hashable, MetaTensor]: + d = dict(data) + for key in self.key_iterator(d): + d[key] = self.padder.inverse(d[key]) + return d + + +class SpatialPadd(Padd): """ Dictionary-based wrapper of :py:class:`monai.transforms.SpatialPad`. Performs padding to the data, symmetric for all sides or all on one side for each dimension. - """ - backend = SpatialPad.backend + """ def __init__( self, keys: KeysCollection, spatial_size: Union[Sequence[int], int], - method: Union[Method, str] = Method.SYMMETRIC, - mode: PadModeSequence = NumpyPadMode.CONSTANT, + method: str = Method.SYMMETRIC, + mode: SequenceStr = PytorchPadMode.CONSTANT, allow_missing_keys: bool = False, **kwargs, ) -> None: @@ -130,7 +175,7 @@ def __init__( keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` spatial_size: the spatial size of output data after padding, if a dimension of the input - data size is bigger than the pad size, will not pad that dimension. + data size is larger than the pad size, will not pad that dimension. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [30, 30, 30] and `spatial_size=[32, 25, -1]`, the spatial size of output data will be [32, 30, 30]. @@ -148,39 +193,11 @@ def __init__( note that `np.pad` treats channel dimension as the first dimension. """ - super().__init__(keys, allow_missing_keys) - self.mode = ensure_tuple_rep(mode, len(self.keys)) - self.padder = SpatialPad(spatial_size, method, **kwargs) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - for key, m in self.key_iterator(d, self.mode): - self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) - d[key] = self.padder(d[key], mode=m) - return d - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = transform[TraceKeys.ORIG_SIZE] - if self.padder.method == Method.SYMMETRIC: - current_size = d[key].shape[1:] - roi_center = [floor(i / 2) if r % 2 == 0 else (i - 1) // 2 for r, i in zip(orig_size, current_size)] - else: - roi_center = [floor(r / 2) if r % 2 == 0 else (r - 1) // 2 for r in orig_size] - - inverse_transform = SpatialCrop(roi_center, orig_size) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d + padder = SpatialPad(spatial_size, method, **kwargs) + super().__init__(keys, padder=padder, mode=mode, allow_missing_keys=allow_missing_keys) -class BorderPadd(MapTransform, InvertibleTransform): +class BorderPadd(Padd): """ Pad the input data by adding specified borders to every dimension. Dictionary-based wrapper of :py:class:`monai.transforms.BorderPad`. @@ -192,7 +209,7 @@ def __init__( self, keys: KeysCollection, spatial_border: Union[Sequence[int], int], - mode: PadModeSequence = NumpyPadMode.CONSTANT, + mode: SequenceStr = PytorchPadMode.CONSTANT, allow_missing_keys: bool = False, **kwargs, ) -> None: @@ -223,43 +240,11 @@ def __init__( note that `np.pad` treats channel dimension as the first dimension. """ - super().__init__(keys, allow_missing_keys) - self.mode = ensure_tuple_rep(mode, len(self.keys)) - self.padder = BorderPad(spatial_border=spatial_border, **kwargs) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - for key, m in self.key_iterator(d, self.mode): - self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) - d[key] = self.padder(d[key], mode=m) - return d + padder = BorderPad(spatial_border=spatial_border, **kwargs) + super().__init__(keys, padder=padder, mode=mode, allow_missing_keys=allow_missing_keys) - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - roi_start = np.array(self.padder.spatial_border) - # Need to convert single value to [min1,min2,...] - if roi_start.size == 1: - roi_start = np.full((len(orig_size)), roi_start) - # need to convert [min1,max1,min2,...] to [min1,min2,...] - elif roi_start.size == 2 * orig_size.size: - roi_start = roi_start[::2] - roi_end = np.array(transform[TraceKeys.ORIG_SIZE]) + roi_start - - inverse_transform = SpatialCrop(roi_start=roi_start, roi_end=roi_end) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d - - -class DivisiblePadd(MapTransform, InvertibleTransform): +class DivisiblePadd(Padd): """ Pad the input data, so that the spatial sizes are divisible by `k`. Dictionary-based wrapper of :py:class:`monai.transforms.DivisiblePad`. @@ -271,8 +256,8 @@ def __init__( self, keys: KeysCollection, k: Union[Sequence[int], int], - mode: PadModeSequence = NumpyPadMode.CONSTANT, - method: Union[Method, str] = Method.SYMMETRIC, + mode: SequenceStr = PytorchPadMode.CONSTANT, + method: str = Method.SYMMETRIC, allow_missing_keys: bool = False, **kwargs, ) -> None: @@ -287,8 +272,8 @@ def __init__( ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. One of the listed string values or a user supplied function. Defaults to ``"constant"``. - See also: https://numpy.org/doc/1.18/reference/ghttps://pytorch.orgenerated/numpy.pad.html - /docs/stable/generated/torch.nn.functional.pad.html + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html It also can be a sequence of string, each element corresponds to a key in ``keys``. method: {``"symmetric"``, ``"end"``} Pad image symmetrically on every side or only pad at the end sides. Defaults to ``"symmetric"``. @@ -299,41 +284,85 @@ def __init__( See also :py:class:`monai.transforms.SpatialPad` """ + padder = DivisiblePad(k=k, method=method, **kwargs) + super().__init__(keys, padder=padder, mode=mode, allow_missing_keys=allow_missing_keys) + + +class Cropd(MapTransform, InvertibleTransform): + """ + Dictionary-based wrapper of abstract class :py:class:`monai.transforms.Crop`. + + Args: + keys: keys of the corresponding items to be transformed. + See also: :py:class:`monai.transforms.compose.MapTransform` + cropper: crop transform for the input image. + allow_missing_keys: don't raise exception if key is missing. + + """ + + backend = Crop.backend + + def __init__(self, keys: KeysCollection, cropper: Crop, allow_missing_keys: bool = False): super().__init__(keys, allow_missing_keys) - self.mode = ensure_tuple_rep(mode, len(self.keys)) - self.padder = DivisiblePad(k=k, method=method, **kwargs) + self.cropper = cropper - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) - for key, m in self.key_iterator(d, self.mode): - self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) - d[key] = self.padder(d[key], mode=m) + for key in self.key_iterator(d): + d[key] = self.cropper(d[key]) # type: ignore return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - + def inverse(self, data: Mapping[Hashable, MetaTensor]) -> Dict[Hashable, MetaTensor]: + d = dict(data) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - current_size = np.array(d[key].shape[1:]) - roi_start = np.floor((current_size - orig_size) / 2) - roi_end = orig_size + roi_start - inverse_transform = SpatialCrop(roi_start=roi_start, roi_end=roi_end) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) + d[key] = self.cropper.inverse(d[key]) + return d + + +class RandCropd(Cropd, Randomizable): + """ + Base class for random crop transform. + Args: + keys: keys of the corresponding items to be transformed. + See also: :py:class:`monai.transforms.compose.MapTransform` + cropper: random crop transform for the input image. + allow_missing_keys: don't raise exception if key is missing. + + """ + + backend = Crop.backend + + def __init__(self, keys: KeysCollection, cropper: Crop, allow_missing_keys: bool = False): + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandCropd": + super().set_random_state(seed, state) + if isinstance(self.cropper, Randomizable): + self.cropper.set_random_state(seed, state) + return self + + def randomize(self, img_size: Sequence[int]) -> None: + if isinstance(self.cropper, Randomizable): + self.cropper.randomize(img_size) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) + # the first key must exist to execute random operations + self.randomize(d[self.first_key(d)].shape[1:]) + for key in self.key_iterator(d): + kwargs = {"randomize": False} if isinstance(self.cropper, Randomizable) else {} + d[key] = self.cropper(d[key], **kwargs) # type: ignore return d -class SpatialCropd(MapTransform, InvertibleTransform): +class SpatialCropd(Cropd): """ Dictionary-based wrapper of :py:class:`monai.transforms.SpatialCrop`. General purpose cropper to produce sub-volume region of interest (ROI). - If a dimension of the expected ROI size is bigger than the input image size, will not crop that dimension. + If a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. It can support to crop ND spatial (channel-first) data. @@ -344,8 +373,6 @@ class SpatialCropd(MapTransform, InvertibleTransform): - the start and end coordinates of the ROI """ - backend = SpatialCrop.backend - def __init__( self, keys: KeysCollection, @@ -361,50 +388,23 @@ def __init__( keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` roi_center: voxel coordinates for center of the crop ROI. - roi_size: size of the crop ROI, if a dimension of ROI size is bigger than image size, + roi_size: size of the crop ROI, if a dimension of ROI size is larger than image size, will not crop that dimension of the image. roi_start: voxel coordinates for start of the crop ROI. roi_end: voxel coordinates for end of the crop ROI, if a coordinate is out of image, use the end coordinate of image. roi_slices: list of slices for each of the spatial dimensions. allow_missing_keys: don't raise exception if key is missing. - """ - super().__init__(keys, allow_missing_keys) - self.cropper = SpatialCrop(roi_center, roi_size, roi_start, roi_end, roi_slices) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - for key in self.key_iterator(d): - self.push_transform(d, key) - d[key] = self.cropper(d[key]) - return d - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - current_size = np.array(d[key].shape[1:]) - # get required pad to start and end - pad_to_start = np.array([s.indices(o)[0] for s, o in zip(self.cropper.slices, orig_size)]) - pad_to_end = orig_size - current_size - pad_to_start - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - return d + """ + cropper = SpatialCrop(roi_center, roi_size, roi_start, roi_end, roi_slices) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) -class CenterSpatialCropd(MapTransform, InvertibleTransform): +class CenterSpatialCropd(Cropd): """ Dictionary-based wrapper of :py:class:`monai.transforms.CenterSpatialCrop`. - If a dimension of the expected ROI size is bigger than the input image size, will not crop that dimension. + If a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. @@ -412,52 +412,21 @@ class CenterSpatialCropd(MapTransform, InvertibleTransform): keys: keys of the corresponding items to be transformed. See also: monai.transforms.MapTransform roi_size: the size of the crop region e.g. [224,224,128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. allow_missing_keys: don't raise exception if key is missing. """ - backend = CenterSpatialCrop.backend - def __init__( self, keys: KeysCollection, roi_size: Union[Sequence[int], int], allow_missing_keys: bool = False ) -> None: - super().__init__(keys, allow_missing_keys) - self.cropper = CenterSpatialCrop(roi_size) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - for key in self.key_iterator(d): - orig_size = d[key].shape[1:] - d[key] = self.cropper(d[key]) - self.push_transform(d, key, orig_size=orig_size) - return d - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - current_size = np.array(d[key].shape[1:]) - pad_to_start = np.floor((orig_size - current_size) / 2).astype(int) - # in each direction, if original size is even and current size is odd, += 1 - pad_to_start[np.logical_and(orig_size % 2 == 0, current_size % 2 == 1)] += 1 - pad_to_end = orig_size - current_size - pad_to_start - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d + cropper = CenterSpatialCrop(roi_size) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) -class CenterScaleCropd(MapTransform, InvertibleTransform): +class CenterScaleCropd(Cropd): """ Dictionary-based wrapper of :py:class:`monai.transforms.CenterScaleCrop`. Note: as using the same scaled ROI to crop, all the input data specified by `keys` should have @@ -471,61 +440,21 @@ class CenterScaleCropd(MapTransform, InvertibleTransform): allow_missing_keys: don't raise exception if key is missing. """ - backend = CenterSpatialCrop.backend - def __init__( self, keys: KeysCollection, roi_scale: Union[Sequence[float], float], allow_missing_keys: bool = False ) -> None: - super().__init__(keys, allow_missing_keys=allow_missing_keys) - self.roi_scale = roi_scale - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - first_key: Union[Hashable, List] = self.first_key(d) - if first_key == []: - return d - - # use the spatial size of first image to scale, expect all images have the same spatial size - img_size = d[first_key].shape[1:] # type: ignore - ndim = len(img_size) - roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.roi_scale, ndim), img_size)] - cropper = CenterSpatialCrop(roi_size) - for key in self.key_iterator(d): - self.push_transform(d, key, orig_size=img_size) - d[key] = cropper(d[key]) - - return d + cropper = CenterScaleCrop(roi_scale) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - current_size = np.array(d[key].shape[1:]) - pad_to_start = np.floor((orig_size - current_size) / 2).astype(int) - # in each direction, if original size is even and current size is odd, += 1 - pad_to_start[np.logical_and(orig_size % 2 == 0, current_size % 2 == 1)] += 1 - pad_to_end = orig_size - current_size - pad_to_start - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d - - -class RandSpatialCropd(Randomizable, MapTransform, InvertibleTransform): +class RandSpatialCropd(RandCropd): """ Dictionary-based version :py:class:`monai.transforms.RandSpatialCrop`. Crop image with random size or specific size ROI. It can crop at a random position as center or at the image center. And allows to set the minimum and maximum size to limit the randomly generated ROI. Suppose all the expected fields specified by `keys` have same shape. - Note: even `random_size=False`, if a dimension of the expected ROI size is bigger than the input image size, + Note: even `random_size=False`, if a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. @@ -534,7 +463,7 @@ class RandSpatialCropd(Randomizable, MapTransform, InvertibleTransform): See also: monai.transforms.MapTransform roi_size: if `random_size` is True, it specifies the minimum crop region. if `random_size` is False, it specifies the expected ROI size to crop. e.g. [224, 224, 128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -548,8 +477,6 @@ class RandSpatialCropd(Randomizable, MapTransform, InvertibleTransform): allow_missing_keys: don't raise exception if key is missing. """ - backend = CenterSpatialCrop.backend - def __init__( self, keys: KeysCollection, @@ -559,78 +486,11 @@ def __init__( random_size: bool = True, allow_missing_keys: bool = False, ) -> None: - MapTransform.__init__(self, keys, allow_missing_keys) - self.roi_size = roi_size - self.max_roi_size = max_roi_size - self.random_center = random_center - self.random_size = random_size - self._slices: Optional[Tuple[slice, ...]] = None - self._size: Optional[Sequence[int]] = None + cropper = RandSpatialCrop(roi_size, max_roi_size, random_center, random_size) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) - def randomize(self, img_size: Sequence[int]) -> None: - self._size = fall_back_tuple(self.roi_size, img_size) - if self.random_size: - max_size = img_size if self.max_roi_size is None else fall_back_tuple(self.max_roi_size, img_size) - if any(i > j for i, j in zip(self._size, max_size)): - raise ValueError(f"min ROI size: {self._size} is bigger than max ROI size: {max_size}.") - self._size = [self.R.randint(low=self._size[i], high=max_size[i] + 1) for i in range(len(img_size))] - if self.random_center: - valid_size = get_valid_patch_size(img_size, self._size) - self._slices = (slice(None),) + get_random_patch(img_size, valid_size, self.R) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - first_key: Union[Hashable, List] = self.first_key(d) - if first_key == []: - return d - - self.randomize(d[first_key].shape[1:]) # type: ignore - if self._size is None: - raise RuntimeError("self._size not specified.") - for key in self.key_iterator(d): - if self.random_center: - self.push_transform(d, key, {"slices": [(i.start, i.stop) for i in self._slices[1:]]}) # type: ignore - d[key] = d[key][self._slices] - else: - self.push_transform(d, key) - cropper = CenterSpatialCrop(self._size) - d[key] = cropper(d[key]) - return d - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = transform[TraceKeys.ORIG_SIZE] - random_center = self.random_center - pad_to_start = np.empty((len(orig_size)), dtype=np.int32) - pad_to_end = np.empty((len(orig_size)), dtype=np.int32) - if random_center: - for i, _slice in enumerate(transform[TraceKeys.EXTRA_INFO]["slices"]): - pad_to_start[i] = _slice[0] - pad_to_end[i] = orig_size[i] - _slice[1] - else: - current_size = d[key].shape[1:] - for i, (o_s, c_s) in enumerate(zip(orig_size, current_size)): - pad_to_start[i] = pad_to_end[i] = (o_s - c_s) / 2 - if o_s % 2 == 0 and c_s % 2 == 1: - pad_to_start[i] += 1 - elif o_s % 2 == 1 and c_s % 2 == 0: - pad_to_end[i] += 1 - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d - - -class RandScaleCropd(RandSpatialCropd): +class RandScaleCropd(RandCropd): """ Dictionary-based version :py:class:`monai.transforms.RandScaleCrop`. Crop image with random size or specific size ROI. @@ -655,8 +515,6 @@ class RandScaleCropd(RandSpatialCropd): allow_missing_keys: don't raise exception if key is missing. """ - backend = RandSpatialCropd.backend - def __init__( self, keys: KeysCollection, @@ -666,50 +524,20 @@ def __init__( random_size: bool = True, allow_missing_keys: bool = False, ) -> None: - super().__init__( - keys=keys, - roi_size=-1, - max_roi_size=None, - random_center=random_center, - random_size=random_size, - allow_missing_keys=allow_missing_keys, - ) - self.roi_scale = roi_scale - self.max_roi_scale = max_roi_scale - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - first_key: Union[Hashable, List] = self.first_key(data) # type: ignore - if first_key == []: - return data # type: ignore - - img_size = data[first_key].shape[1:] # type: ignore - ndim = len(img_size) - self.roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.roi_scale, ndim), img_size)] - if self.max_roi_scale is not None: - self.max_roi_size = [ceil(r * s) for r, s in zip(ensure_tuple_rep(self.max_roi_scale, ndim), img_size)] - else: - self.max_roi_size = None - return super().__call__(data=data) - - -@contextlib.contextmanager -def _nullcontext(x): - """ - This is just like contextlib.nullcontext but also works in Python 3.6. - """ - yield x + cropper = RandScaleCrop(roi_scale, max_roi_scale, random_center, random_size) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) -class RandSpatialCropSamplesd(Randomizable, MapTransform, InvertibleTransform): +class RandSpatialCropSamplesd(Randomizable, MapTransform): """ Dictionary-based version :py:class:`monai.transforms.RandSpatialCropSamples`. Crop image with random size or specific size ROI to generate a list of N samples. It can crop at a random position as center or at the image center. And allows to set the minimum size to limit the randomly generated ROI. Suppose all the expected fields - specified by `keys` have same shape, and add `patch_index` to the corresponding meta data. + specified by `keys` have same shape, and add `patch_index` to the corresponding metadata. It will return a list of dictionaries for all the cropped images. - Note: even `random_size=False`, if a dimension of the expected ROI size is bigger than the input image size, + Note: even `random_size=False`, if a dimension of the expected ROI size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected ROI, and the cropped results of several images may not have exactly the same shape. @@ -718,7 +546,7 @@ class RandSpatialCropSamplesd(Randomizable, MapTransform, InvertibleTransform): See also: monai.transforms.MapTransform roi_size: if `random_size` is True, it specifies the minimum crop region. if `random_size` is False, it specifies the expected ROI size to crop. e.g. [224, 224, 128] - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. If its components have non-positive values, the corresponding size of input image will be used. for example: if the spatial size of input data is [40, 40, 40] and `roi_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -729,15 +557,6 @@ class RandSpatialCropSamplesd(Randomizable, MapTransform, InvertibleTransform): random_center: crop at random position as center or the image center. random_size: crop with random size or specific size ROI. The actual size is sampled from `randint(roi_size, img_size)`. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - used to add `patch_index` to the meta dict. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - used to add `patch_index` to the meta dict. allow_missing_keys: don't raise exception if key is missing. Raises: @@ -745,8 +564,10 @@ class RandSpatialCropSamplesd(Randomizable, MapTransform, InvertibleTransform): """ - backend = RandSpatialCropd.backend + backend = RandSpatialCropSamples.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") def __init__( self, keys: KeysCollection, @@ -756,63 +577,33 @@ def __init__( random_center: bool = True, random_size: bool = True, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_missing_keys: bool = False, ) -> None: MapTransform.__init__(self, keys, allow_missing_keys) - if num_samples < 1: - raise ValueError(f"num_samples must be positive, got {num_samples}.") - self.num_samples = num_samples - self.cropper = RandSpatialCropd(keys, roi_size, max_roi_size, random_center, random_size, allow_missing_keys) - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - - def set_random_state( - self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None - ) -> "RandSpatialCropSamplesd": - super().set_random_state(seed, state) - self.cropper.set_random_state(seed, state) - return self + self.cropper = RandSpatialCropSamples(roi_size, num_samples, max_roi_size, random_center, random_size) def randomize(self, data: Optional[Any] = None) -> None: - pass - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict[Hashable, NdarrayOrTensor]]: - ret = [] - for i in range(self.num_samples): - d = dict(data) - # deep copy all the unmodified data - for key in set(data.keys()).difference(set(self.keys)): - d[key] = deepcopy(data[key]) - cropped = self.cropper(d) - # self.cropper will have added RandSpatialCropd to the list. Change to RandSpatialCropSamplesd - for key in self.key_iterator(cropped): - cropped[self.trace_key(key)][-1][TraceKeys.CLASS_NAME] = self.__class__.__name__ # type: ignore - cropped[self.trace_key(key)][-1][TraceKeys.ID] = id(self) # type: ignore - # add `patch_index` to the meta data - for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - if meta_key not in cropped: - cropped[meta_key] = {} # type: ignore - cropped[meta_key][Key.PATCH_INDEX] = i # type: ignore - ret.append(cropped) + self.sub_seed = self.R.randint(MAX_SEED, dtype="uint32") + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> List[Dict[Hashable, torch.Tensor]]: + # output starts as empty list of dictionaries + ret: List[Dict[Hashable, torch.Tensor]] = [{} for _ in range(self.cropper.num_samples)] + # deep copy all the unmodified data + for key in set(data.keys()).difference(set(self.keys)): + for r in ret: + r[key] = deepcopy(data[key]) + + # for each key we reset the random state to ensure crops are the same + self.randomize() + for key in self.key_iterator(dict(data)): + self.cropper.set_random_state(seed=self.sub_seed) + for i, im in enumerate(self.cropper(data[key])): + ret[i][key] = im return ret - def inverse(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]: - d = deepcopy(dict(data)) - # We changed the transform name from RandSpatialCropd to RandSpatialCropSamplesd - # Need to revert that since we're calling RandSpatialCropd's inverse - for key in self.key_iterator(d): - d[self.trace_key(key)][-1][TraceKeys.CLASS_NAME] = self.cropper.__class__.__name__ - d[self.trace_key(key)][-1][TraceKeys.ID] = id(self.cropper) - context_manager = allow_missing_keys_mode if self.allow_missing_keys else _nullcontext - with context_manager(self.cropper): - return self.cropper.inverse(d) - -class CropForegroundd(MapTransform, InvertibleTransform): +class CropForegroundd(Cropd): """ Dictionary-based version :py:class:`monai.transforms.CropForeground`. Crop only the foreground object of the expected images. @@ -825,8 +616,6 @@ class CropForegroundd(MapTransform, InvertibleTransform): channels. And it can also add margin to every dim of the bounding box of foreground object. """ - backend = CropForeground.backend - def __init__( self, keys: KeysCollection, @@ -836,11 +625,11 @@ def __init__( margin: Union[Sequence[int], int] = 0, allow_smaller: bool = True, k_divisible: Union[Sequence[int], int] = 1, - mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = NumpyPadMode.CONSTANT, + mode: SequenceStr = PytorchPadMode.CONSTANT, start_coord_key: str = "foreground_start_coord", end_coord_key: str = "foreground_end_coord", allow_missing_keys: bool = False, - **np_kwargs, + **pad_kwargs, ) -> None: """ Args: @@ -852,7 +641,7 @@ def __init__( of image. if None, select foreground on the whole image. margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims. allow_smaller: when computing box size with `margin`, whether allow the image size to be smaller - than box size, default to `True`. if the margined size is bigger than image size, will pad with + than box size, default to `True`. if the margined size is larger than image size, will pad with specified `mode`. k_divisible: make each spatial dimension to be divisible by k, default to 1. if `k_divisible` is an int, the same `k` be applied to all the input spatial dimensions. @@ -866,64 +655,38 @@ def __init__( start_coord_key: key to record the start coordinate of spatial bounding box for foreground. end_coord_key: key to record the end coordinate of spatial bounding box for foreground. allow_missing_keys: don't raise exception if key is missing. - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ - super().__init__(keys, allow_missing_keys) self.source_key = source_key self.start_coord_key = start_coord_key self.end_coord_key = end_coord_key - self.cropper = CropForeground( + cropper = CropForeground( select_fn=select_fn, channel_indices=channel_indices, margin=margin, allow_smaller=allow_smaller, k_divisible=k_divisible, - **np_kwargs, + **pad_kwargs, ) + super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys) self.mode = ensure_tuple_rep(mode, len(self.keys)) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) + self.cropper: CropForeground box_start, box_end = self.cropper.compute_bounding_box(img=d[self.source_key]) - d[self.start_coord_key] = box_start - d[self.end_coord_key] = box_end + if self.start_coord_key is not None: + d[self.start_coord_key] = box_start + if self.end_coord_key is not None: + d[self.end_coord_key] = box_end for key, m in self.key_iterator(d, self.mode): - self.push_transform(d, key, extra_info={"box_start": box_start, "box_end": box_end}) d[key] = self.cropper.crop_pad(img=d[key], box_start=box_start, box_end=box_end, mode=m) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.asarray(transform[TraceKeys.ORIG_SIZE]) - cur_size = np.asarray(d[key].shape[1:]) - extra_info = transform[TraceKeys.EXTRA_INFO] - box_start = np.asarray(extra_info["box_start"]) - box_end = np.asarray(extra_info["box_end"]) - # first crop the padding part - roi_start = np.maximum(-box_start, 0) - roi_end = cur_size - np.maximum(box_end - orig_size, 0) - - d[key] = SpatialCrop(roi_start=roi_start, roi_end=roi_end)(d[key]) - - # update bounding box to pad - pad_to_start = np.maximum(box_start, 0) - pad_to_end = orig_size - np.minimum(box_end, orig_size) - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - # second pad back the original size - d[key] = BorderPad(pad)(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - - return d - -class RandWeightedCropd(Randomizable, MapTransform, InvertibleTransform): +class RandWeightedCropd(Randomizable, MapTransform): """ Samples a list of `num_samples` image patches according to the provided `weight_map`. @@ -935,16 +698,6 @@ class RandWeightedCropd(Randomizable, MapTransform, InvertibleTransform): 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. - center_coord_key: if specified, the actual sampling location will be stored with the corresponding key. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - used to add `patch_index` to the meta dict. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - used to add `patch_index` to the meta dict. allow_missing_keys: don't raise exception if key is missing. See Also: @@ -953,6 +706,9 @@ class RandWeightedCropd(Randomizable, MapTransform, InvertibleTransform): backend = SpatialCrop.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") + @deprecated_arg(name="center_coord_key", since="0.9", msg_suffix="coords stored in img.meta['crop_center']") def __init__( self, keys: KeysCollection, @@ -961,103 +717,59 @@ def __init__( num_samples: int = 1, center_coord_key: Optional[str] = None, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_missing_keys: bool = False, ): MapTransform.__init__(self, keys, allow_missing_keys) - self.spatial_size = ensure_tuple(spatial_size) self.w_key = w_key - self.num_samples = int(num_samples) - self.center_coord_key = center_coord_key - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - self.centers: List[np.ndarray] = [] - - def randomize(self, weight_map: NdarrayOrTensor) -> None: - self.centers = weighted_patch_samples( - spatial_size=self.spatial_size, w=weight_map[0], n_samples=self.num_samples, r_state=self.R - ) + self.cropper = RandWeightedCrop(spatial_size, num_samples) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict[Hashable, NdarrayOrTensor]]: - d = dict(data) - self.randomize(d[self.w_key]) - _spatial_size = fall_back_tuple(self.spatial_size, d[self.w_key].shape[1:]) - - # initialize returned list with shallow copy to preserve key ordering - results: List[Dict[Hashable, NdarrayOrTensor]] = [dict(data) for _ in range(self.num_samples)] - # fill in the extra keys with unmodified data - for i in range(self.num_samples): - for key in set(data.keys()).difference(set(self.keys)): - results[i][key] = deepcopy(data[key]) - for key in self.key_iterator(d): - img = d[key] - if img.shape[1:] != d[self.w_key].shape[1:]: - raise ValueError( - f"data {key} and weight map {self.w_key} spatial shape mismatch: " - f"{img.shape[1:]} vs {d[self.w_key].shape[1:]}." - ) - for i, center in enumerate(self.centers): - cropper = SpatialCrop(roi_center=center, roi_size=_spatial_size) - orig_size = img.shape[1:] - results[i][key] = cropper(img) - self.push_transform(results[i], key, extra_info={"center": center}, orig_size=orig_size) - if self.center_coord_key: - results[i][self.center_coord_key] = center - # fill in the extra keys with unmodified data - for i in range(self.num_samples): - # add `patch_index` to the meta data - for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - if meta_key not in results[i]: - results[i][meta_key] = {} # type: ignore - results[i][meta_key][Key.PATCH_INDEX] = i # type: ignore - - return results - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.asarray(transform[TraceKeys.ORIG_SIZE]) - current_size = np.asarray(d[key].shape[1:]) - center = transform[TraceKeys.EXTRA_INFO]["center"] - cropper = SpatialCrop(roi_center=center, roi_size=self.spatial_size) - # get required pad to start and end - pad_to_start = np.array([s.indices(o)[0] for s, o in zip(cropper.slices, orig_size)]) - pad_to_end = orig_size - current_size - pad_to_start - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandWeightedCropd": + super().set_random_state(seed, state) + self.cropper.set_random_state(seed, state) + return self - return d + def randomize(self, weight_map: NdarrayOrTensor) -> None: + self.cropper.randomize(weight_map) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> List[Dict[Hashable, torch.Tensor]]: + # output starts as empty list of dictionaries + ret: List = [{} for _ in range(self.cropper.num_samples)] + # deep copy all the unmodified data + for key in set(data.keys()).difference(set(self.keys)): + for r in ret: + r[key] = deepcopy(data[key]) + + self.randomize(weight_map=data[self.w_key]) + for key in self.key_iterator(data): + for i, im in enumerate(self.cropper(data[key], weight_map=data[self.w_key], randomize=False)): + ret[i][key] = im + return ret -class RandCropByPosNegLabeld(Randomizable, MapTransform, InvertibleTransform): +class RandCropByPosNegLabeld(Randomizable, MapTransform): """ Dictionary-based version :py:class:`monai.transforms.RandCropByPosNegLabel`. Crop random fixed sized regions with the center being a foreground or background voxel based on the Pos Neg Ratio. Suppose all the expected fields specified by `keys` have same shape, - and add `patch_index` to the corresponding meta data. + and add `patch_index` to the corresponding metadata. And will return a list of dictionaries for all the cropped images. - If a dimension of the expected spatial size is bigger than the input image size, + If a dimension of the expected spatial size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than the expected size, and the cropped results of several images may not have exactly the same shape. + And if the crop ROI is partly out of the image, will automatically adjust the crop center + to ensure the valid crop ROI. Args: keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` label_key: name of key for label image, this will be used for finding foreground/background. spatial_size: the spatial size of the crop region e.g. [224, 224, 128]. - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. if its components have non-positive values, the corresponding size of `data[label_key]` will be used. for example: if the spatial size of input data is [40, 40, 40] and `spatial_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -1078,15 +790,6 @@ class RandCropByPosNegLabeld(Randomizable, MapTransform, InvertibleTransform): `image_threshold`, and randomly select crop centers based on them, need to provide `fg_indices_key` and `bg_indices_key` together, expect to be 1 dim array of spatial indices after flattening. a typical usage is to call `FgBgToIndicesd` transform first and cache the results. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - used to add `patch_index` to the meta dict. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - used to add `patch_index` to the meta dict. allow_smaller: if `False`, an exception will be raised if the image is smaller than the requested ROI in any dimension. If `True`, any smaller dimensions will be set to match the cropped size (i.e., no cropping in that dimension). @@ -1100,6 +803,8 @@ class RandCropByPosNegLabeld(Randomizable, MapTransform, InvertibleTransform): backend = RandCropByPosNegLabel.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") def __init__( self, keys: KeysCollection, @@ -1113,55 +818,41 @@ def __init__( fg_indices_key: Optional[str] = None, bg_indices_key: Optional[str] = None, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_smaller: bool = False, allow_missing_keys: bool = False, ) -> None: MapTransform.__init__(self, keys, allow_missing_keys) self.label_key = label_key - self.spatial_size: Union[Tuple[int, ...], Sequence[int], int] = spatial_size - if pos < 0 or neg < 0: - raise ValueError(f"pos and neg must be nonnegative, got pos={pos} neg={neg}.") - if pos + neg == 0: - raise ValueError("Incompatible values: pos=0 and neg=0.") - self.pos_ratio = pos / (pos + neg) - self.num_samples = num_samples self.image_key = image_key - self.image_threshold = image_threshold self.fg_indices_key = fg_indices_key self.bg_indices_key = bg_indices_key - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - self.centers: Optional[List[List[int]]] = None - self.allow_smaller = allow_smaller + self.cropper = RandCropByPosNegLabel( + spatial_size=spatial_size, + pos=pos, + neg=neg, + num_samples=num_samples, + image_threshold=image_threshold, + allow_smaller=allow_smaller, + ) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandCropByPosNegLabeld": + super().set_random_state(seed, state) + self.cropper.set_random_state(seed, state) + return self def randomize( self, - label: NdarrayOrTensor, + label: torch.Tensor, fg_indices: Optional[NdarrayOrTensor] = None, bg_indices: Optional[NdarrayOrTensor] = None, - image: Optional[NdarrayOrTensor] = None, + image: Optional[torch.Tensor] = None, ) -> None: - self.spatial_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) - if fg_indices is None or bg_indices is None: - fg_indices_, bg_indices_ = map_binary_to_indices(label, image, self.image_threshold) - else: - fg_indices_ = fg_indices - bg_indices_ = bg_indices - self.centers = generate_pos_neg_label_crop_centers( - self.spatial_size, - self.num_samples, - self.pos_ratio, - label.shape[1:], - fg_indices_, - bg_indices_, - self.R, - self.allow_smaller, - ) + self.cropper.randomize(label=label, fg_indices=fg_indices, bg_indices=bg_indices, image=image) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict[Hashable, NdarrayOrTensor]]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> List[Dict[Hashable, torch.Tensor]]: d = dict(data) label = d[self.label_key] image = d[self.image_key] if self.image_key else None @@ -1169,57 +860,21 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict[Hashab bg_indices = d.pop(self.bg_indices_key, None) if self.bg_indices_key is not None else None self.randomize(label, fg_indices, bg_indices, image) - if not isinstance(self.spatial_size, tuple): - raise ValueError("spatial_size must be a valid tuple.") - if self.centers is None: - raise ValueError("no available ROI centers to crop.") # initialize returned list with shallow copy to preserve key ordering - results: List[Dict[Hashable, NdarrayOrTensor]] = [dict(d) for _ in range(self.num_samples)] - - for i, center in enumerate(self.centers): - # fill in the extra keys with unmodified data - for key in set(d.keys()).difference(set(self.keys)): - results[i][key] = deepcopy(d[key]) - for key in self.key_iterator(d): - img = d[key] - cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) - orig_size = img.shape[1:] - results[i][key] = cropper(img) - self.push_transform(results[i], key, extra_info={"center": center}, orig_size=orig_size) - # add `patch_index` to the meta data - for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - if meta_key not in results[i]: - results[i][meta_key] = {} # type: ignore - results[i][meta_key][Key.PATCH_INDEX] = i # type: ignore - - return results - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.asarray(transform[TraceKeys.ORIG_SIZE]) - current_size = np.asarray(d[key].shape[1:]) - center = transform[TraceKeys.EXTRA_INFO]["center"] - cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) # type: ignore - # get required pad to start and end - pad_to_start = np.array([s.indices(o)[0] for s, o in zip(cropper.slices, orig_size)]) - pad_to_end = orig_size - current_size - pad_to_start - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) + ret: List = [{} for _ in range(self.cropper.num_samples)] + # deep copy all the unmodified data + for key in set(d.keys()).difference(set(self.keys)): + for r in ret: + r[key] = deepcopy(d[key]) - return d + for key in self.key_iterator(d): + for i, im in enumerate(self.cropper(d[key], label=label, randomize=False)): + ret[i][key] = im + return ret -class RandCropByLabelClassesd(Randomizable, MapTransform, InvertibleTransform): +class RandCropByLabelClassesd(Randomizable, MapTransform): """ Dictionary-based version :py:class:`monai.transforms.RandCropByLabelClasses`. Crop random fixed sized regions with the center being a class based on the specified ratios of every class. @@ -1257,16 +912,18 @@ class RandCropByLabelClassesd(Randomizable, MapTransform, InvertibleTransform): [0, 1, 3], [1, 2, 1], [0, 0, 0]] [1, 3, 0]] - If a dimension of the expected spatial size is bigger than the input image size, + If a dimension of the expected spatial size is larger than the input image size, will not crop that dimension. So the cropped result may be smaller than expected size, and the cropped results of several images may not have exactly same shape. + And if the crop ROI is partly out of the image, will automatically adjust the crop center to ensure the + valid crop ROI. Args: keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` label_key: name of key for label image, this will be used for finding indices of every class. spatial_size: the spatial size of the crop region e.g. [224, 224, 128]. - if a dimension of ROI size is bigger than image size, will not crop that dimension of the image. + if a dimension of ROI size is larger than image size, will not crop that dimension of the image. if its components have non-positive values, the corresponding size of `label` will be used. for example: if the spatial size of input data is [40, 40, 40] and `spatial_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40]. @@ -1282,15 +939,6 @@ class RandCropByLabelClassesd(Randomizable, MapTransform, InvertibleTransform): `image_threshold`, and randomly select crop centers based on them, expect to be 1 dim array of spatial indices after flattening. a typical usage is to call `ClassesToIndices` transform first and cache the results for better performance. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - used to add `patch_index` to the meta dict. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - used to add `patch_index` to the meta dict. allow_smaller: if `False`, an exception will be raised if the image is smaller than the requested ROI in any dimension. If `True`, any smaller dimensions will remain unchanged. @@ -1300,6 +948,8 @@ class RandCropByLabelClassesd(Randomizable, MapTransform, InvertibleTransform): backend = RandCropByLabelClasses.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") def __init__( self, keys: KeysCollection, @@ -1312,99 +962,57 @@ def __init__( image_threshold: float = 0.0, indices_key: Optional[str] = None, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_smaller: bool = False, allow_missing_keys: bool = False, ) -> None: MapTransform.__init__(self, keys, allow_missing_keys) self.label_key = label_key - self.spatial_size: Union[Tuple[int, ...], Sequence[int], int] = spatial_size - self.ratios = ratios - self.num_classes = num_classes - self.num_samples = num_samples self.image_key = image_key - self.image_threshold = image_threshold self.indices_key = indices_key - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - self.centers: Optional[List[List[int]]] = None - self.allow_smaller = allow_smaller + self.cropper = RandCropByLabelClasses( + spatial_size=spatial_size, + ratios=ratios, + num_classes=num_classes, + num_samples=num_samples, + image_threshold=image_threshold, + allow_smaller=allow_smaller, + ) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandCropByLabelClassesd": + super().set_random_state(seed, state) + self.cropper.set_random_state(seed, state) + return self def randomize( - self, - label: NdarrayOrTensor, - indices: Optional[List[NdarrayOrTensor]] = None, - image: Optional[NdarrayOrTensor] = None, + self, label: torch.Tensor, indices: Optional[List[NdarrayOrTensor]] = None, image: Optional[torch.Tensor] = None ) -> None: - self.spatial_size = fall_back_tuple(self.spatial_size, default=label.shape[1:]) - if indices is None: - indices_ = map_classes_to_indices(label, self.num_classes, image, self.image_threshold) - else: - indices_ = indices - self.centers = generate_label_classes_crop_centers( - self.spatial_size, self.num_samples, label.shape[1:], indices_, self.ratios, self.R, self.allow_smaller - ) + self.cropper.randomize(label=label, indices=indices, image=image) - def __call__(self, data: Mapping[Hashable, Any]) -> List[Dict[Hashable, NdarrayOrTensor]]: + def __call__(self, data: Mapping[Hashable, Any]) -> List[Dict[Hashable, torch.Tensor]]: d = dict(data) label = d[self.label_key] image = d[self.image_key] if self.image_key else None indices = d.pop(self.indices_key, None) if self.indices_key is not None else None self.randomize(label, indices, image) - if not isinstance(self.spatial_size, tuple): - raise ValueError("spatial_size must be a valid tuple.") - if self.centers is None: - raise ValueError("no available ROI centers to crop.") # initialize returned list with shallow copy to preserve key ordering - results: List[Dict[Hashable, NdarrayOrTensor]] = [dict(d) for _ in range(self.num_samples)] - - for i, center in enumerate(self.centers): - # fill in the extra keys with unmodified data - for key in set(d.keys()).difference(set(self.keys)): - results[i][key] = deepcopy(d[key]) - for key in self.key_iterator(d): - img = d[key] - cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) - orig_size = img.shape[1:] - results[i][key] = cropper(img) - self.push_transform(results[i], key, extra_info={"center": center}, orig_size=orig_size) - # add `patch_index` to the meta data - for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - if meta_key not in results[i]: - results[i][meta_key] = {} # type: ignore - results[i][meta_key][Key.PATCH_INDEX] = i # type: ignore - - return results - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.asarray(transform[TraceKeys.ORIG_SIZE]) - current_size = np.asarray(d[key].shape[1:]) - center = transform[TraceKeys.EXTRA_INFO]["center"] - cropper = SpatialCrop(roi_center=tuple(center), roi_size=self.spatial_size) # type: ignore - # get required pad to start and end - pad_to_start = np.array([s.indices(o)[0] for s, o in zip(cropper.slices, orig_size)]) - pad_to_end = orig_size - current_size - pad_to_start - # interleave mins and maxes - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - inverse_transform = BorderPad(pad) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) + ret: List = [{} for _ in range(self.cropper.num_samples)] + # deep copy all the unmodified data + for key in set(d.keys()).difference(set(self.keys)): + for r in ret: + r[key] = deepcopy(d[key]) - return d + for key in self.key_iterator(d): + for i, im in enumerate(self.cropper(d[key], label=label, randomize=False)): + ret[i][key] = im + return ret -class ResizeWithPadOrCropd(MapTransform, InvertibleTransform): +class ResizeWithPadOrCropd(Padd): """ Dictionary-based wrapper of :py:class:`monai.transforms.ResizeWithPadOrCrop`. @@ -1413,76 +1021,32 @@ class ResizeWithPadOrCropd(MapTransform, InvertibleTransform): See also: monai.transforms.MapTransform spatial_size: the spatial size of output data after padding or crop. If has non-positive values, the corresponding size of input image will be used (no padding). - mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, - ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} - One of the listed string values or a user supplied function for padding. Defaults to ``"constant"``. + mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, + ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. + One of the listed string values or a user supplied function. Defaults to ``"constant"``. See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html It also can be a sequence of string, each element corresponds to a key in ``keys``. allow_missing_keys: don't raise exception if key is missing. method: {``"symmetric"``, ``"end"``} Pad image symmetrically on every side or only pad at the end sides. Defaults to ``"symmetric"``. - np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension. - more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + note that `np.pad` treats channel dimension as the first dimension. """ - backend = ResizeWithPadOrCrop.backend - def __init__( self, keys: KeysCollection, spatial_size: Union[Sequence[int], int], - mode: NumpyPadModeSequence = NumpyPadMode.CONSTANT, + mode: SequenceStr = PytorchPadMode.CONSTANT, allow_missing_keys: bool = False, - method: Union[Method, str] = Method.SYMMETRIC, - **np_kwargs, + method: str = Method.SYMMETRIC, + **pad_kwargs, ) -> None: - super().__init__(keys, allow_missing_keys) - self.mode = ensure_tuple_rep(mode, len(self.keys)) - self.padcropper = ResizeWithPadOrCrop(spatial_size=spatial_size, method=method, **np_kwargs) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = dict(data) - for key, m in self.key_iterator(d, self.mode): - orig_size = d[key].shape[1:] - d[key] = self.padcropper(d[key], mode=m) - self.push_transform(d, key, orig_size=orig_size, extra_info={"mode": m.value if isinstance(m, Enum) else m}) - return d - - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_size = np.array(transform[TraceKeys.ORIG_SIZE]) - current_size = np.array(d[key].shape[1:]) - # Unfortunately, we can't just use ResizeWithPadOrCrop with original size because of odd/even rounding. - # Instead, we first pad any smaller dimensions, and then we crop any larger dimensions. - - # First, do pad - if np.any((orig_size - current_size) > 0): - pad_to_start = np.floor((orig_size - current_size) / 2).astype(int) - # in each direction, if original size is even and current size is odd, += 1 - pad_to_start[np.logical_and(orig_size % 2 == 0, current_size % 2 == 1)] += 1 - pad_to_start[pad_to_start < 0] = 0 - pad_to_end = orig_size - current_size - pad_to_start - pad_to_end[pad_to_end < 0] = 0 - pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist()))) - d[key] = BorderPad(pad)(d[key]) - - # Next crop - if np.any((orig_size - current_size) < 0): - if self.padcropper.padder.method == Method.SYMMETRIC: - roi_center = [floor(i / 2) if r % 2 == 0 else (i - 1) // 2 for r, i in zip(orig_size, current_size)] - else: - roi_center = [floor(r / 2) if r % 2 == 0 else (r - 1) // 2 for r in orig_size] - - d[key] = SpatialCrop(roi_center, orig_size)(d[key]) - - # Remove the applied transform - self.pop_transform(d, key) - - return d + padcropper = ResizeWithPadOrCrop(spatial_size=spatial_size, method=method, **pad_kwargs) + super().__init__(keys, padder=padcropper, mode=mode, allow_missing_keys=allow_missing_keys) # type: ignore class BoundingRectd(MapTransform): @@ -1525,9 +1089,12 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d +PadD = PadDict = Padd SpatialPadD = SpatialPadDict = SpatialPadd BorderPadD = BorderPadDict = BorderPadd DivisiblePadD = DivisiblePadDict = DivisiblePadd +CropD = CropDict = Cropd +RandCropD = RandCropDict = RandCropd SpatialCropD = SpatialCropDict = SpatialCropd CenterSpatialCropD = CenterSpatialCropDict = CenterSpatialCropd CenterScaleCropD = CenterScaleCropDict = CenterScaleCropd diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index da46b105e1..3fa3aa63fb 100644 --- a/monai/transforms/intensity/array.py +++ b/monai/transforms/intensity/array.py @@ -16,7 +16,7 @@ from abc import abstractmethod from collections.abc import Iterable from functools import partial -from typing import Any, Callable, List, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union from warnings import warn import numpy as np @@ -24,24 +24,19 @@ 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.utils import get_random_patch, get_valid_patch_size from monai.networks.layers import GaussianFilter, HilbertTransform, SavitzkyGolayFilter from monai.transforms.transform import RandomizableTransform, Transform from monai.transforms.utils import Fourier, equalize_hist, is_positive, rescale_array from monai.transforms.utils_pytorch_numpy_unification import clip, percentile, where -from monai.utils import ( - InvalidPyTorchVersionError, - convert_data_type, - convert_to_dst_type, - ensure_tuple, - ensure_tuple_rep, - ensure_tuple_size, - fall_back_tuple, - pytorch_after, -) from monai.utils.deprecate_utils import deprecated_arg from monai.utils.enums import TransformBackends -from monai.utils.type_conversion import convert_to_tensor, get_equivalent_dtype +from monai.utils.misc import ensure_tuple, ensure_tuple_rep, ensure_tuple_size, fall_back_tuple +from monai.utils.module import min_version, optional_import +from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor, get_equivalent_dtype + +skimage, _ = optional_import("skimage", "0.19.0", min_version) __all__ = [ "RandGaussianNoise", @@ -77,6 +72,7 @@ "HistogramNormalize", "IntensityRemap", "RandIntensityRemap", + "ForegroundMask", ] @@ -114,6 +110,7 @@ def __call__(self, img: NdarrayOrTensor, mean: Optional[float] = None, randomize """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize(img=img, mean=self.mean if mean is None else mean) @@ -192,6 +189,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: super().randomize(None) @@ -234,6 +232,7 @@ def __call__(self, img: NdarrayOrTensor, offset: Optional[float] = None) -> Ndar Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) offset = self.offset if offset is None else offset out = img + offset out, *_ = convert_data_type(data=out, dtype=img.dtype) @@ -263,7 +262,7 @@ def __init__(self, offsets: Union[Tuple[float, float], float], prob: float = 0.1 else: self.offsets = (min(offsets), max(offsets)) self._offset = self.offsets[0] - self._shfiter = ShiftIntensity(self._offset) + self._shifter = ShiftIntensity(self._offset) def randomize(self, data: Optional[Any] = None) -> None: super().randomize(None) @@ -281,13 +280,14 @@ def __call__(self, img: NdarrayOrTensor, factor: Optional[float] = None, randomi can be some image specific value at runtime, like: max(img), etc. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() if not self._do_transform: return img - return self._shfiter(img, self._offset if factor is None else self._offset * factor) + return self._shifter(img, self._offset if factor is None else self._offset * factor) class StdShiftIntensity(Transform): @@ -335,6 +335,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if self.dtype is not None: img, *_ = convert_data_type(img, dtype=self.dtype) if self.channel_wise: @@ -393,6 +394,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -445,17 +447,18 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: ValueError: When ``self.minv=None`` or ``self.maxv=None`` and ``self.factor=None``. Incompatible values. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_t = convert_to_tensor(img, track_meta=False) ret: NdarrayOrTensor if self.minv is not None or self.maxv is not None: if self.channel_wise: - out = [rescale_array(d, self.minv, self.maxv, dtype=self.dtype) for d in img] - ret = torch.stack(out) if isinstance(img, torch.Tensor) else np.stack(out) # type: ignore + out = [rescale_array(d, self.minv, self.maxv, dtype=self.dtype) for d in img_t] + ret = torch.stack(out) # type: ignore else: - ret = rescale_array(img, self.minv, self.maxv, dtype=self.dtype) + ret = rescale_array(img_t, self.minv, self.maxv, dtype=self.dtype) else: - ret = (img * (1 + self.factor)) if self.factor is not None else img - - ret, *_ = convert_data_type(ret, dtype=self.dtype or img.dtype) + ret = (img_t * (1 + self.factor)) if self.factor is not None else img_t + ret = convert_to_dst_type(ret, dst=img, dtype=self.dtype or img_t.dtype)[0] return ret @@ -498,6 +501,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -579,6 +583,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize(img_size=img.shape[1:]) @@ -601,7 +606,8 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen class NormalizeIntensity(Transform): """ - Normalize input based on provided args, using calculated mean and std if not provided. + Normalize input based on the `subtrahend` and `divisor`: `(img - subtrahend) / divisor`. + Use calculated mean or std value of the input image if no `subtrahend` or `divisor` provided. This transform can normalize only non-zero values or entire image, and can also calculate mean and std on each channel separately. When `channel_wise` is True, the first dimension of `subtrahend` and `divisor` should @@ -680,6 +686,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`, assuming `img` is a channel-first array if `self.channel_wise` is True, """ + img = convert_to_tensor(img, track_meta=get_track_meta()) dtype = self.dtype or img.dtype if self.channel_wise: if self.subtrahend is not None and len(self.subtrahend) != len(img): @@ -724,6 +731,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) mask = img > self.threshold if self.above else img < self.threshold res = where(mask, img, self.cval) res, *_ = convert_data_type(res, dtype=img.dtype) @@ -735,7 +743,7 @@ class ScaleIntensityRange(Transform): Apply specific intensity scaling to the whole numpy array. Scaling from [a_min, a_max] to [b_min, b_max] with clip option. - When `b_min` or `b_max` are `None`, `scacled_array * (b_max - b_min) + b_min` will be skipped. + When `b_min` or `b_max` are `None`, `scaled_array * (b_max - b_min) + b_min` will be skipped. If `clip=True`, when `b_min`/`b_max` is None, the clipping is not performed on the corresponding edge. Args: @@ -769,6 +777,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) dtype = self.dtype or img.dtype if self.a_max - self.a_min == 0.0: warn("Divide by zero (a_min == a_max)", Warning) @@ -807,6 +816,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) epsilon = 1e-7 img_min = img.min() img_range = img.max() - img_min @@ -854,6 +864,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -969,19 +980,21 @@ def _normalize(self, img: NdarrayOrTensor) -> NdarrayOrTensor: a_min=a_min, a_max=a_max, b_min=b_min, b_max=b_max, clip=self.clip, dtype=self.dtype ) img = scalar(img) + img = convert_to_tensor(img, track_meta=False) return img def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_t = convert_to_tensor(img, track_meta=False) if self.channel_wise: - out = [self._normalize(img=d) for d in img] - img = torch.stack(out) if isinstance(img, torch.Tensor) else np.stack(out) # type: ignore + img_t = torch.stack([self._normalize(img=d) for d in img_t]) # type: ignore else: - img = self._normalize(img=img) + img_t = self._normalize(img=img_t) - return img + return convert_to_dst_type(img_t, dst=img)[0] class MaskIntensity(Transform): @@ -1021,6 +1034,7 @@ def __call__(self, img: NdarrayOrTensor, mask_data: Optional[NdarrayOrTensor] = - ValueError: When ``mask_data`` and ``img`` channels differ and ``mask_data`` is not single channel. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) mask_data = self.mask_data if mask_data is None else mask_data if mask_data is None: raise ValueError("must provide the mask_data when initializing the transform or at runtime.") @@ -1034,7 +1048,7 @@ def __call__(self, img: NdarrayOrTensor, mask_data: Optional[NdarrayOrTensor] = f"got img channels={img.shape[0]} mask_data channels={mask_data_.shape[0]}." ) - return img * mask_data_ + return convert_to_dst_type(img * mask_data_, dst=img)[0] class SavitzkyGolaySmooth(Transform): @@ -1071,7 +1085,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: array containing smoothed result. """ - self.img_t = convert_to_tensor(img) + img = convert_to_tensor(img, track_meta=get_track_meta()) + self.img_t = convert_to_tensor(img, track_meta=False) # add one to transform axis because a batch axis will be added at dimension 0 savgol_filter = SavitzkyGolayFilter(self.window_length, self.order, self.axis + 1, self.mode) @@ -1085,7 +1100,6 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: class DetectEnvelope(Transform): """ Find the envelope of the input data along the requested axis using a Hilbert transform. - Requires PyTorch 1.7.0+ and the PyTorch FFT module (which is not included in NVIDIA PyTorch Release 20.10). Args: axis: Axis along which to detect the envelope. Default 1, i.e. the first spatial dimension. @@ -1098,9 +1112,6 @@ class DetectEnvelope(Transform): def __init__(self, axis: int = 1, n: Union[int, None] = None) -> None: - if not pytorch_after(1, 7): - raise InvalidPyTorchVersionError("1.7.0", self.__class__.__name__) - if axis < 0: raise ValueError("axis must be zero or positive.") @@ -1117,6 +1128,7 @@ def __call__(self, img: NdarrayOrTensor): np.ndarray containing envelope of data in img along the specified axis. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) img_t, *_ = convert_data_type(img, torch.Tensor) # add one to transform axis because a batch axis will be added at dimension 0 hilbert_transform = HilbertTransform(self.axis + 1, self.n) @@ -1148,6 +1160,7 @@ def __init__(self, sigma: Union[Sequence[float], float] = 1.0, approx: str = "er self.approx = approx 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) sigma: Union[Sequence[torch.Tensor], torch.Tensor] if isinstance(self.sigma, Sequence): @@ -1156,7 +1169,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor: sigma = torch.as_tensor(self.sigma, device=img_t.device) gaussian_filter = GaussianFilter(img_t.ndim - 1, sigma, approx=self.approx) out_t: torch.Tensor = gaussian_filter(img_t.unsqueeze(0)).squeeze(0) - out, *_ = convert_data_type(out_t, type(img), device=img.device if isinstance(img, torch.Tensor) else None) + out, *_ = convert_to_dst_type(out_t, dst=img, dtype=out_t.dtype) return out @@ -1204,6 +1217,7 @@ def randomize(self, data: Optional[Any] = None) -> None: self.z = self.R.uniform(low=self.sigma_z[0], high=self.sigma_z[1]) def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -1256,6 +1270,7 @@ def __init__( self.approx = approx 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.float32) gf1, gf2 = ( @@ -1265,7 +1280,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor: blurred_f = gf1(img_t.unsqueeze(0)) filter_blurred_f = gf2(blurred_f) out_t: torch.Tensor = (blurred_f + self.alpha * (blurred_f - filter_blurred_f)).squeeze(0) - out, *_ = convert_data_type(out_t, type(img), device=img.device if isinstance(img, torch.Tensor) else None) + out, *_ = convert_to_dst_type(out_t, dst=img, dtype=out_t.dtype) return out @@ -1338,6 +1353,7 @@ def randomize(self, data: Optional[Any] = None) -> None: self.a = self.R.uniform(low=self.alpha[0], high=self.alpha[1]) def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -1362,7 +1378,7 @@ class RandHistogramShift(RandomizableTransform): prob: probability of histogram shift. """ - backend = [TransformBackends.NUMPY] + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] def __init__(self, num_control_points: Union[Tuple[int, int], int] = 10, prob: float = 0.1) -> None: RandomizableTransform.__init__(self, prob) @@ -1377,8 +1393,25 @@ def __init__(self, num_control_points: Union[Tuple[int, int], int] = 10, prob: f if min(num_control_points) <= 2: raise ValueError("num_control_points should be greater than or equal to 3") self.num_control_points = (min(num_control_points), max(num_control_points)) - self.reference_control_points: np.ndarray - self.floating_control_points: np.ndarray + self.reference_control_points: NdarrayOrTensor + self.floating_control_points: NdarrayOrTensor + + def interp(self, x: NdarrayOrTensor, xp: NdarrayOrTensor, fp: NdarrayOrTensor) -> NdarrayOrTensor: + ns = torch if isinstance(x, torch.Tensor) else np + if isinstance(x, np.ndarray): + # approx 2x faster than code below for ndarray + return np.interp(x, xp, fp) + + m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1]) + b = fp[:-1] - (m * xp[:-1]) + + indices = ns.searchsorted(xp.reshape(-1), x.reshape(-1)) - 1 + indices = ns.clip(indices, 0, len(m) - 1) + + f = (m[indices] * x.reshape(-1) + b[indices]).reshape(x.shape) + f[x < xp[0]] = fp[0] + f[x > xp[-1]] = fp[-1] + return f def randomize(self, data: Optional[Any] = None) -> None: super().randomize(None) @@ -1393,6 +1426,7 @@ def randomize(self, data: Optional[Any] = None) -> None: ) def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -1401,15 +1435,14 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen if self.reference_control_points is None or self.floating_control_points is None: raise RuntimeError("please call the `randomize()` function first.") - img_np, *_ = convert_data_type(img, np.ndarray) - img_min, img_max = img_np.min(), img_np.max() - reference_control_points_scaled = self.reference_control_points * (img_max - img_min) + img_min - floating_control_points_scaled = self.floating_control_points * (img_max - img_min) + img_min - img_np = np.asarray( # type: ignore - np.interp(img_np, reference_control_points_scaled, floating_control_points_scaled), dtype=img_np.dtype - ) - img, *_ = convert_to_dst_type(img_np, dst=img) - return img + img_t = convert_to_tensor(img, track_meta=False) + xp, *_ = convert_to_dst_type(self.reference_control_points, dst=img_t) + yp, *_ = convert_to_dst_type(self.floating_control_points, dst=img_t) + img_min, img_max = img_t.min(), img_t.max() + reference_control_points_scaled = xp * (img_max - img_min) + img_min + floating_control_points_scaled = yp * (img_max - img_min) + img_min + img_t = self.interp(img_t, reference_control_points_scaled, floating_control_points_scaled) + return convert_to_dst_type(img_t, dst=img)[0] class GibbsNoise(Transform, Fourier): @@ -1442,14 +1475,17 @@ def __init__(self, alpha: float = 0.1, as_tensor_output: bool = True) -> None: self.alpha = alpha def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: - n_dims = len(img.shape[1:]) + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_t = convert_to_tensor(img, track_meta=False) + n_dims = len(img_t.shape[1:]) # FT - k = self.shift_fourier(img, n_dims) + k = self.shift_fourier(img_t, n_dims) # build and apply mask k = self._apply_mask(k) # map back - img = self.inv_shift_fourier(k, n_dims) + out = self.inv_shift_fourier(k, n_dims) + img, *_ = convert_to_dst_type(out, dst=img, dtype=out.dtype) return img @@ -1535,6 +1571,7 @@ def randomize(self, data: Any) -> None: self.sampled_alpha = self.R.uniform(self.alpha[0], self.alpha[1]) def __call__(self, img: NdarrayOrTensor, randomize: bool = True): + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: # randomize application and possibly alpha self.randomize(None) @@ -1609,6 +1646,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Args: img: image with dimensions (C, H, W) or (C, H, W, D) """ + img = convert_to_tensor(img, track_meta=get_track_meta()) # checking that tuples in loc are consistent with img size self._check_indices(img) @@ -1751,7 +1789,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True): raise RuntimeError( "If intensity_range is a sequence of sequences, then there must be one (low, high) tuple for each channel." ) - + img = convert_to_tensor(img, track_meta=get_track_meta()) self.sampled_k_intensity = [] self.sampled_locs = [] @@ -1886,6 +1924,7 @@ def _transform_holes(self, img: np.ndarray) -> np.ndarray: raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize(img.shape[1:]) @@ -2047,6 +2086,7 @@ def __init__( self.dtype = dtype def __call__(self, img: NdarrayOrTensor, mask: Optional[NdarrayOrTensor] = None) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) img_np, *_ = convert_data_type(img, np.ndarray) mask = mask if mask is not None else self.mask mask_np: Optional[np.ndarray] = None @@ -2081,10 +2121,6 @@ class IntensityRemap(RandomizableTransform): curve. slope: slope of the linear component. Easiest to leave default value and tune the kernel_size parameter instead. - return_map: set to True for the transform to return a dictionary version - of the lookup table used in the intensity remapping. The keys - correspond to the old intensities, and the values are the new - values. """ def __init__(self, kernel_size: int = 30, slope: float = 0.7): @@ -2099,10 +2135,10 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: Args: img: image to remap. """ - - img = img.clone() + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_ = convert_to_tensor(img, track_meta=False) # sample noise - vals_to_sample = torch.unique(img).tolist() + vals_to_sample = torch.unique(img_).tolist() noise = torch.from_numpy(self.R.choice(vals_to_sample, len(vals_to_sample) - 1 + self.kernel_size)) # smooth noise = torch.nn.AvgPool1d(self.kernel_size, stride=1)(noise.unsqueeze(0)).squeeze() @@ -2110,11 +2146,11 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: grid = torch.arange(len(noise)) / len(noise) noise += self.slope * grid # rescale - noise = (noise - noise.min()) / (noise.max() - noise.min()) * img.max() + img.min() + noise = (noise - noise.min()) / (noise.max() - noise.min()) * img_.max() + img_.min() # intensity remapping function - index_img = torch.bucketize(img, torch.tensor(vals_to_sample)) - img = noise[index_img] + index_img = torch.bucketize(img_, torch.tensor(vals_to_sample)) + img, *_ = convert_to_dst_type(noise[index_img], dst=img) return img @@ -2147,7 +2183,7 @@ def __init__(self, prob: float = 0.1, kernel_size: int = 30, slope: float = 0.7, RandomizableTransform.__init__(self, prob=prob) self.kernel_size = kernel_size self.slope = slope - self.channel_wise = True + self.channel_wise = channel_wise def __call__(self, img: torch.Tensor) -> torch.Tensor: """ @@ -2155,6 +2191,7 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: img: image to remap. """ super().randomize(None) + img = convert_to_tensor(img, track_meta=get_track_meta()) if self._do_transform: if self.channel_wise: img = torch.stack( @@ -2167,3 +2204,100 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: img = IntensityRemap(self.kernel_size, self.R.choice([-self.slope, self.slope]))(img) return img + + +class ForegroundMask(Transform): + """ + Creates a binary mask that defines the foreground based on thresholds in RGB or HSV color space. + This transform receives an RGB (or grayscale) image where by default it is assumed that the foreground has + low values (dark) while the background has high values (white). Otherwise, set `invert` argument to `True`. + + Args: + threshold: an int or a float number that defines the threshold that values less than that are foreground. + It also can be a callable that receives each dimension of the image and calculate the threshold, + or a string that defines such callable from `skimage.filter.threshold_...`. For the list of available + threshold functions, please refer to https://scikit-image.org/docs/stable/api/skimage.filters.html + Moreover, a dictionary can be passed that defines such thresholds for each channel, like + {"R": 100, "G": "otsu", "B": skimage.filter.threshold_mean} + hsv_threshold: similar to threshold but HSV color space ("H", "S", and "V"). + Unlike RBG, in HSV, value greater than `hsv_threshold` are considered foreground. + invert: invert the intensity range of the input image, so that the dtype maximum is now the dtype minimum, + and vice-versa. + + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + threshold: Union[Dict, Callable, str, float, int] = "otsu", + hsv_threshold: Optional[Union[Dict, Callable, str, float, int]] = None, + invert: bool = False, + ) -> None: + self.thresholds: Dict[str, Union[Callable, float]] = {} + if threshold is not None: + if isinstance(threshold, dict): + for mode, th in threshold.items(): + self._set_threshold(th, mode.upper()) + else: + self._set_threshold(threshold, "R") + self._set_threshold(threshold, "G") + self._set_threshold(threshold, "B") + if hsv_threshold is not None: + if isinstance(hsv_threshold, dict): + for mode, th in hsv_threshold.items(): + self._set_threshold(th, mode.upper()) + else: + self._set_threshold(hsv_threshold, "H") + self._set_threshold(hsv_threshold, "S") + self._set_threshold(hsv_threshold, "V") + + self.thresholds = {k: v for k, v in self.thresholds.items() if v is not None} + if self.thresholds.keys().isdisjoint(set("RGBHSV")): + raise ValueError( + f"Threshold for at least one channel of RGB or HSV needs to be set. {self.thresholds} is provided." + ) + self.invert = invert + + def _set_threshold(self, threshold, mode): + if callable(threshold): + self.thresholds[mode] = threshold + elif isinstance(threshold, str): + self.thresholds[mode] = getattr(skimage.filters, "threshold_" + threshold.lower()) + elif isinstance(threshold, (float, int)): + self.thresholds[mode] = float(threshold) + else: + raise ValueError( + f"`threshold` should be either a callable, string, or float number, {type(threshold)} was given." + ) + + def _get_threshold(self, image, mode): + threshold = self.thresholds.get(mode) + if callable(threshold): + return threshold(image) + return threshold + + def __call__(self, image: NdarrayOrTensor): + image = convert_to_tensor(image, track_meta=get_track_meta()) + img_rgb, *_ = convert_data_type(image, np.ndarray) + if self.invert: + img_rgb = skimage.util.invert(img_rgb) + foregrounds = [] + if not self.thresholds.keys().isdisjoint(set("RGB")): + rgb_foreground = np.zeros_like(img_rgb[:1]) + for img, mode in zip(img_rgb, "RGB"): + threshold = self._get_threshold(img, mode) + if threshold: + rgb_foreground = np.logical_or(rgb_foreground, img <= threshold) + foregrounds.append(rgb_foreground) + if not self.thresholds.keys().isdisjoint(set("HSV")): + img_hsv = skimage.color.rgb2hsv(img_rgb, channel_axis=0) + hsv_foreground = np.zeros_like(img_rgb[:1]) + for img, mode in zip(img_hsv, "HSV"): + threshold = self._get_threshold(img, mode) + if threshold: + hsv_foreground = np.logical_or(hsv_foreground, img > threshold) + foregrounds.append(hsv_foreground) + + mask = np.stack(foregrounds).all(axis=0) + return convert_to_dst_type(src=mask, dst=image)[0] diff --git a/monai/transforms/intensity/dictionary.py b/monai/transforms/intensity/dictionary.py index b0f5149456..b9308255cf 100644 --- a/monai/transforms/intensity/dictionary.py +++ b/monai/transforms/intensity/dictionary.py @@ -21,8 +21,10 @@ from monai.config import DtypeLike, KeysCollection from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_obj import get_track_meta from monai.transforms.intensity.array import ( AdjustContrast, + ForegroundMask, GaussianSharpen, GaussianSmooth, GibbsNoise, @@ -54,7 +56,7 @@ ) from monai.transforms.transform import MapTransform, RandomizableTransform from monai.transforms.utils import is_positive -from monai.utils import ensure_tuple, ensure_tuple_rep +from monai.utils import convert_to_tensor, ensure_tuple, ensure_tuple_rep from monai.utils.deprecate_utils import deprecated_arg from monai.utils.enums import PostFix @@ -88,6 +90,7 @@ "RandCoarseDropoutd", "RandCoarseShuffled", "HistogramNormalized", + "ForegroundMaskd", "RandGaussianNoiseD", "RandGaussianNoiseDict", "ShiftIntensityD", @@ -146,6 +149,8 @@ "HistogramNormalizeDict", "RandKSpaceSpikeNoiseD", "RandKSpaceSpikeNoiseDict", + "ForegroundMaskD", + "ForegroundMaskDict", ] DEFAULT_POST_FIX = PostFix.meta() @@ -193,11 +198,15 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random noise first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d self.rand_gaussian_noise.randomize(d[first_key]) # type: ignore @@ -268,6 +277,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d for key in self.key_iterator(d): @@ -297,18 +308,18 @@ def __init__( See also: :py:class:`monai.transforms.compose.MapTransform` offset: offset value to shift the intensity of image. factor_key: if not None, use it as the key to extract a value from the corresponding - meta data dictionary of `key` at runtime, and multiply the `offset` to shift intensity. + metadata dictionary of `key` at runtime, and multiply the `offset` to shift intensity. Usually, `IntensityStatsd` transform can pre-compute statistics of intensity values - and store in the meta data. + and store in the metadata. it also can be a sequence of strings, map to `keys`. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. used to extract the factor value is `factor_key` is not None. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. used to extract the factor value is `factor_key` is not None. allow_missing_keys: don't raise exception if key is missing. """ @@ -356,18 +367,18 @@ def __init__( offsets: offset range to randomly shift. if single number, offset value is picked from (-offsets, offsets). factor_key: if not None, use it as the key to extract a value from the corresponding - meta data dictionary of `key` at runtime, and multiply the random `offset` to shift intensity. + metadata dictionary of `key` at runtime, and multiply the random `offset` to shift intensity. Usually, `IntensityStatsd` transform can pre-compute statistics of intensity values - and store in the meta data. + and store in the metadata. it also can be a sequence of strings, map to `keys`. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. used to extract the factor value is `factor_key` is not None. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. used to extract the factor value is `factor_key` is not None. prob: probability of rotating. (Default 0.1, with 10% probability it returns a rotated array.) @@ -394,6 +405,8 @@ def __call__(self, data) -> Dict[Hashable, NdarrayOrTensor]: d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random shift factor @@ -490,6 +503,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random shift factor @@ -584,6 +599,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random scale factor @@ -637,11 +654,15 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random bias factor first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d self.rand_bias_field.randomize(img_size=d[first_key].shape[1:]) # type: ignore @@ -829,6 +850,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random gamma value @@ -1042,6 +1065,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random sigma @@ -1157,6 +1182,8 @@ def __call__(self, data: Dict[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Ndar d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random sigma1, sigma2, etc. @@ -1205,6 +1232,8 @@ def __call__(self, data: Dict[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Ndar d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random shift params @@ -1266,6 +1295,8 @@ def __call__(self, data: Dict[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Ndar d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # all the keys share the same random noise params @@ -1450,6 +1481,8 @@ def __call__(self, data: Dict[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Ndar d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d for key in self.key_iterator(d): @@ -1526,11 +1559,15 @@ def __call__(self, data): d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # expect all the specified keys have same spatial shape and share same random holes first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d self.dropper.randomize(d[first_key].shape[1:]) @@ -1595,11 +1632,15 @@ def __call__(self, data): d = dict(data) self.randomize(None) if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d # expect all the specified keys have same spatial shape and share same random holes first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) return d self.shuffle.randomize(d[first_key].shape[1:]) @@ -1654,6 +1695,52 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d +class ForegroundMaskd(MapTransform): + """ + Creates a binary mask that defines the foreground based on thresholds in RGB or HSV color space. + This transform receives an RGB (or grayscale) image where by default it is assumed that the foreground has + low values (dark) while the background is white. + + Args: + keys: keys of the corresponding items to be transformed. + threshold: an int or a float number that defines the threshold that values less than that are foreground. + It also can be a callable that receives each dimension of the image and calculate the threshold, + or a string that defines such callable from `skimage.filter.threshold_...`. For the list of available + threshold functions, please refer to https://scikit-image.org/docs/stable/api/skimage.filters.html + Moreover, a dictionary can be passed that defines such thresholds for each channel, like + {"R": 100, "G": "otsu", "B": skimage.filter.threshold_mean} + hsv_threshold: similar to threshold but HSV color space ("H", "S", and "V"). + Unlike RBG, in HSV, value greater than `hsv_threshold` are considered foreground. + invert: invert the intensity range of the input image, so that the dtype maximum is now the dtype minimum, + and vice-versa. + new_key_prefix: this prefix be prepended to the key to create a new key for the output and keep the value of + key intact. By default not prefix is set and the corresponding array to the key will be replaced. + allow_missing_keys: do not raise exception if key is missing. + + """ + + def __init__( + self, + keys: KeysCollection, + threshold: Union[Dict, Callable, str, float] = "otsu", + hsv_threshold: Optional[Union[Dict, Callable, str, float, int]] = None, + invert: bool = False, + new_key_prefix: Optional[str] = None, + allow_missing_keys: bool = False, + ) -> None: + super().__init__(keys, allow_missing_keys) + self.transform = ForegroundMask(threshold=threshold, hsv_threshold=hsv_threshold, invert=invert) + self.new_key_prefix = new_key_prefix + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + new_key = key if self.new_key_prefix is None else self.new_key_prefix + key + d[new_key] = self.transform(d[key]) + + return d + + RandGaussianNoiseD = RandGaussianNoiseDict = RandGaussianNoised RandRicianNoiseD = RandRicianNoiseDict = RandRicianNoised ShiftIntensityD = ShiftIntensityDict = ShiftIntensityd @@ -1683,3 +1770,4 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N RandCoarseDropoutD = RandCoarseDropoutDict = RandCoarseDropoutd HistogramNormalizeD = HistogramNormalizeDict = HistogramNormalized RandCoarseShuffleD = RandCoarseShuffleDict = RandCoarseShuffled +ForegroundMaskD = ForegroundMaskDict = ForegroundMaskd diff --git a/monai/transforms/inverse.py b/monai/transforms/inverse.py index c8bfeeca05..41a02989db 100644 --- a/monai/transforms/inverse.py +++ b/monai/transforms/inverse.py @@ -8,11 +8,15 @@ # 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. + import os -from typing import Hashable, Mapping, Optional, Tuple +import warnings +from contextlib import contextmanager +from typing import Any, Hashable, Mapping, Optional, Tuple import torch +from monai.data.meta_tensor import MetaTensor from monai.transforms.transform import Transform from monai.utils.enums import TraceKeys @@ -21,24 +25,37 @@ class TraceableTransform(Transform): """ - Maintains a stack of applied transforms. The stack is inserted as pairs of - `trace_key: list of transforms` to each data dictionary. + Maintains a stack of applied transforms to data. + + Data can be one of two types: + 1. A `MetaTensor` (this is the preferred data type). + 2. A dictionary of data containing arrays/tensors and auxiliary metadata. In + this case, a key must be supplied (this dictionary-based approach is deprecated). + + If `data` is of type `MetaTensor`, then the applied transform will be added to ``data.applied_operations``. + + If `data` is a dictionary, then one of two things can happen: + 1. If data[key] is a `MetaTensor`, the applied transform will be added to ``data[key].applied_operations``. + 2. Else, the applied transform will be appended to an adjacent list using + `trace_key`. If, for example, the key is `image`, then the transform + will be appended to `image_transforms` (this dictionary-based approach is deprecated). + + Hopefully it is clear that there are three total possibilities: + 1. data is `MetaTensor` + 2. data is dictionary, data[key] is `MetaTensor` + 3. data is dictionary, data[key] is not `MetaTensor` (this is a deprecated approach). The ``__call__`` method of this transform class must be implemented so - that the transformation information for each key is stored when - ``__call__`` is called. If the transforms were applied to keys "image" and - "label", there will be two extra keys in the dictionary: "image_transforms" - and "label_transforms" (based on `TraceKeys.KEY_SUFFIX`). Each list - contains a list of the transforms applied to that key. + that the transformation information is stored during the data transformation. - The information in ``data[key_transform]`` will be compatible with the - default collate since it only stores strings, numbers and arrays. + The information in the stack of applied transforms must be compatible with the + default collate, by only storing strings, numbers and arrays. `tracing` could be enabled by `self.set_tracing` or setting `MONAI_TRACE_TRANSFORM` when initializing the class. """ - tracing = False if os.environ.get("MONAI_TRACE_TRANSFORM", "1") == "0" else True + tracing = not os.environ.get("MONAI_TRACE_TRANSFORM", "1") == "0" def set_tracing(self, tracing: bool) -> None: """Set whether to trace transforms.""" @@ -48,37 +65,152 @@ def set_tracing(self, tracing: bool) -> None: def trace_key(key: Hashable = None): """The key to store the stack of applied transforms.""" if key is None: - return TraceKeys.KEY_SUFFIX - return str(key) + TraceKeys.KEY_SUFFIX + return f"{TraceKeys.KEY_SUFFIX}" + return f"{key}{TraceKeys.KEY_SUFFIX}" - def push_transform( - self, data: Mapping, key: Hashable = None, extra_info: Optional[dict] = None, orig_size: Optional[Tuple] = None - ) -> None: - """PUsh to a stack of applied transforms for that key.""" - if not self.tracing: - return + def get_transform_info( + self, data, key: Hashable = None, extra_info: Optional[dict] = None, orig_size: Optional[Tuple] = None + ) -> dict: + """ + Return a dictionary with the relevant information pertaining to an applied transform. + + Args: + data: input data. Can be dictionary or MetaTensor. We can use `shape` to + determine the original size of the object (unless that has been given + explicitly, see `orig_size`). + key: if data is a dictionary, data[key] will be modified. + extra_info: if desired, any extra information pertaining to the applied + transform can be stored in this dictionary. These are often needed for + computing the inverse transformation. + orig_size: sometimes during the inverse it is useful to know what the size + of the original image was, in which case it can be supplied here. + + Returns: + Dictionary of data pertaining to the applied transformation. + """ info = {TraceKeys.CLASS_NAME: self.__class__.__name__, TraceKeys.ID: id(self)} if orig_size is not None: info[TraceKeys.ORIG_SIZE] = orig_size - elif key in data and hasattr(data[key], "shape"): + elif isinstance(data, Mapping) and key in data and hasattr(data[key], "shape"): info[TraceKeys.ORIG_SIZE] = data[key].shape[1:] + elif hasattr(data, "shape"): + info[TraceKeys.ORIG_SIZE] = data.shape[1:] if extra_info is not None: info[TraceKeys.EXTRA_INFO] = extra_info # If class is randomizable transform, store whether the transform was actually performed (based on `prob`) if hasattr(self, "_do_transform"): # RandomizableTransform info[TraceKeys.DO_TRANSFORM] = self._do_transform # type: ignore - # If this is the first, create list - if self.trace_key(key) not in data: - if not isinstance(data, dict): - data = dict(data) - data[self.trace_key(key)] = [] - data[self.trace_key(key)].append(info) - - def pop_transform(self, data: Mapping, key: Hashable = None): - """Remove the most recent applied transform.""" + return info + + def push_transform( + self, data, key: Hashable = None, extra_info: Optional[dict] = None, orig_size: Optional[Tuple] = None + ) -> None: + """ + Push to a stack of applied transforms. + + Args: + data: dictionary of data or `MetaTensor`. + key: if data is a dictionary, data[key] will be modified. + extra_info: if desired, any extra information pertaining to the applied + transform can be stored in this dictionary. These are often needed for + computing the inverse transformation. + orig_size: sometimes during the inverse it is useful to know what the size + of the original image was, in which case it can be supplied here. + + Returns: + None, but data has been updated to store the applied transformation. + """ if not self.tracing: return - return data.get(self.trace_key(key), []).pop() + info = self.get_transform_info(data, key, extra_info, orig_size) + + if isinstance(data, MetaTensor): + data.push_applied_operation(info) + elif isinstance(data, Mapping): + if key in data and isinstance(data[key], MetaTensor): + data[key].push_applied_operation(info) + else: + # If this is the first, create list + if self.trace_key(key) not in data: + if not isinstance(data, dict): + data = dict(data) + data[self.trace_key(key)] = [] + data[self.trace_key(key)].append(info) + else: + warnings.warn(f"`data` should be either `MetaTensor` or dictionary, got {type(data)}. {info} not tracked.") + + def check_transforms_match(self, transform: Mapping) -> None: + """Check transforms are of same instance.""" + xform_id = transform.get(TraceKeys.ID, "") + if xform_id == id(self): + return + # TraceKeys.NONE to skip the id check + if xform_id == TraceKeys.NONE: + return + xform_name = transform.get(TraceKeys.CLASS_NAME, "") + # basic check if multiprocessing uses 'spawn' (objects get recreated so don't have same ID) + if torch.multiprocessing.get_start_method() in ("spawn", None) and xform_name == self.__class__.__name__: + return + raise RuntimeError( + f"Error {self.__class__.__name__} getting the most recently " + f"applied invertible transform {xform_name} {xform_id} != {id(self)}." + ) + + def get_most_recent_transform(self, data, key: Hashable = None, check: bool = True, pop: bool = False): + """ + Get most recent transform for the stack. + + Args: + data: dictionary of data or `MetaTensor`. + key: if data is a dictionary, data[key] will be modified. + check: if true, check that `self` is the same type as the most recently-applied transform. + pop: if true, remove the transform as it is returned. + + Returns: + Dictionary of most recently applied transform + + Raises: + - RuntimeError: data is neither `MetaTensor` nor dictionary + """ + if not self.tracing: + raise RuntimeError("Transform Tracing must be enabled to get the most recent transform.") + if isinstance(data, MetaTensor): + all_transforms = data.applied_operations + elif isinstance(data, Mapping): + if key in data and isinstance(data[key], MetaTensor): + all_transforms = data[key].applied_operations + else: + all_transforms = data.get(self.trace_key(key), MetaTensor.get_default_applied_operations()) + else: + raise ValueError(f"`data` should be either `MetaTensor` or dictionary, got {type(data)}.") + if check: + self.check_transforms_match(all_transforms[-1]) + return all_transforms.pop() if pop else all_transforms[-1] + + def pop_transform(self, data, key: Hashable = None, check: bool = True): + """ + Return and pop the most recent transform. + + Args: + data: dictionary of data or `MetaTensor` + key: if data is a dictionary, data[key] will be modified + check: if true, check that `self` is the same type as the most recently-applied transform. + + Returns: + Dictionary of most recently applied transform + + Raises: + - RuntimeError: data is neither `MetaTensor` nor dictionary + """ + return self.get_most_recent_transform(data, key, check, pop=True) + + @contextmanager + def trace_transform(self, to_trace: bool): + """Temporarily set the tracing status of a transform with a context manager.""" + prev = self.tracing + self.tracing = to_trace + yield + self.tracing = prev class InvertibleTransform(TraceableTransform): @@ -95,7 +227,7 @@ class InvertibleTransform(TraceableTransform): different parameters being passed to each label (e.g., different interpolation for image and label). - - the inverse transforms are applied in a last- in-first-out order. As + - the inverse transforms are applied in a last-in-first-out order. As the inverse is applied, its entry is removed from the list detailing the applied transformations. That is to say that during the forward pass, the list of applied transforms grows, and then during the @@ -118,26 +250,7 @@ class InvertibleTransform(TraceableTransform): """ - def check_transforms_match(self, transform: Mapping) -> None: - """Check transforms are of same instance.""" - xform_name = transform.get(TraceKeys.CLASS_NAME, "") - xform_id = transform.get(TraceKeys.ID, "") - if xform_id == id(self): - return - # basic check if multiprocessing uses 'spawn' (objects get recreated so don't have same ID) - if torch.multiprocessing.get_start_method() in ("spawn", None) and xform_name == self.__class__.__name__: - return - raise RuntimeError(f"Error inverting the most recently applied invertible transform {xform_name} {xform_id}.") - - def get_most_recent_transform(self, data: Mapping, key: Hashable = None): - """Get most recent transform.""" - if not self.tracing: - raise RuntimeError("Transform Tracing must be enabled to get the most recent transform.") - transform = data[self.trace_key(key)][-1] - self.check_transforms_match(transform) - return transform - - def inverse(self, data: dict) -> dict: + def inverse(self, data: Any) -> Any: """ Inverse of ``__call__``. diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py index 5bafd84eaf..8dd849d33e 100644 --- a/monai/transforms/io/array.py +++ b/monai/transforms/io/array.py @@ -20,7 +20,7 @@ import warnings from pathlib import Path from pydoc import locate -from typing import Dict, List, Optional, Sequence, Union +from typing import Dict, List, Optional, Sequence, Type, Union import numpy as np import torch @@ -28,20 +28,32 @@ from monai.config import DtypeLike, NdarrayOrTensor, PathLike from monai.data import image_writer from monai.data.folder_layout import FolderLayout -from monai.data.image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader +from monai.data.image_reader import ( + ImageReader, + ITKReader, + NibabelReader, + NrrdReader, + NumpyReader, + PILReader, + PydicomReader, +) +from monai.data.meta_tensor import MetaTensor from monai.transforms.transform import Transform from monai.transforms.utility.array import EnsureChannelFirst -from monai.utils import GridSampleMode, GridSamplePadMode +from monai.utils import GridSamplePadMode from monai.utils import ImageMetaKey as Key -from monai.utils import InterpolateMode, OptionalImportError, ensure_tuple, look_up_option, optional_import +from monai.utils import OptionalImportError, convert_to_dst_type, ensure_tuple, look_up_option, optional_import nib, _ = optional_import("nibabel") Image, _ = optional_import("PIL.Image") +nrrd, _ = optional_import("nrrd") __all__ = ["LoadImage", "SaveImage", "SUPPORTED_READERS"] SUPPORTED_READERS = { + "pydicomreader": PydicomReader, "itkreader": ITKReader, + "nrrdreader": NrrdReader, "numpyreader": NumpyReader, "pilreader": PILReader, "nibabelreader": NibabelReader, @@ -85,7 +97,7 @@ class LoadImage(Transform): - User-specified reader in the constructor of `LoadImage`. - Readers from the last to the first in the registered list. - Current default readers: (nii, nii.gz -> NibabelReader), (png, jpg, bmp -> PILReader), - (npz, npy -> NumpyReader), (DICOM file -> ITKReader). + (npz, npy -> NumpyReader), (nrrd -> NrrdReader), (DICOM file -> ITKReader). See also: @@ -99,29 +111,31 @@ def __init__( image_only: bool = False, dtype: DtypeLike = np.float32, ensure_channel_first: bool = False, + simple_keys: bool = False, *args, **kwargs, ) -> None: """ Args: - reader: reader to load image file and meta data + reader: reader to load image file and metadata - if `reader` is None, a default set of `SUPPORTED_READERS` will be used. - if `reader` is a string, it's treated as a class name or dotted path (such as ``"monai.data.ITKReader"``), the supported built-in reader classes are - ``"ITKReader"``, ``"NibabelReader"``, ``"NumpyReader"``. + ``"ITKReader"``, ``"NibabelReader"``, ``"NumpyReader"``, ``"PydicomReader"``. a reader instance will be constructed with the `*args` and `**kwargs` parameters. - if `reader` is a reader class/instance, it will be registered to this loader accordingly. - image_only: if True return only the image volume, otherwise return image data array and header dict. + image_only: if True return only the image MetaTensor, otherwise return image and header dict. dtype: if not None convert the loaded image to this data type. - ensure_channel_first: if `True` and loaded both image array and meta data, automatically convert + ensure_channel_first: if `True` and loaded both image array and metadata, automatically convert the image array shape to `channel first`. default to `False`. + simple_keys: whether to remove redundant metadata keys, default to False for backward compatibility. args: additional parameters for reader if providing a reader name. kwargs: additional parameters for reader if providing a reader name. Note: - - The transform returns an image data array if `image_only` is True, - or a tuple of two elements containing the data array, and the meta data in a dictionary format otherwise. + - The transform returns a MetaTensor, unless `set_track_meta(False)` has been used, in which case, a + `torch.Tensor` will be returned. - If `reader` is specified, the loader will attempt to use the specified readers and the default supported readers. This might introduce overheads when handling the exceptions of trying the incompatible loaders. In this case, it is therefore recommended setting the most appropriate reader as @@ -133,6 +147,7 @@ def __init__( self.image_only = image_only self.dtype = dtype self.ensure_channel_first = ensure_channel_first + self.simple_keys = simple_keys self.readers: List[ImageReader] = [] for r in SUPPORTED_READERS: # set predefined readers as default @@ -174,7 +189,7 @@ def __init__( def register(self, reader: ImageReader): """ - Register image reader to load image file and meta data. + Register image reader to load image file and metadata. Args: reader: reader instance to be registered with this loader. @@ -186,7 +201,7 @@ def register(self, reader: ImageReader): def __call__(self, filename: Union[Sequence[PathLike], PathLike], reader: Optional[ImageReader] = None): """ - Load image file and meta data from the given filename(s). + Load image file and metadata from the given filename(s). If `reader` is not specified, this class automatically chooses readers based on the reversed order of registered readers `self.readers`. @@ -197,7 +212,7 @@ def __call__(self, filename: Union[Sequence[PathLike], PathLike], reader: Option and will stack them together as multi-channels data. if provided directory path instead of file path, will treat it as DICOM images series and read. - reader: runtime reader to load image file and meta data. + reader: runtime reader to load image file and metadata. """ filename = tuple(f"{Path(s).expanduser()}" for s in ensure_tuple(filename)) # allow Path objects @@ -236,19 +251,19 @@ def __call__(self, filename: Union[Sequence[PathLike], PathLike], reader: Option img_array: NdarrayOrTensor img_array, meta_data = reader.get_data(img) - img_array = img_array.astype(self.dtype, copy=False) + img_array = convert_to_dst_type(img_array, dst=img_array, dtype=self.dtype)[0] if not isinstance(meta_data, dict): raise ValueError("`meta_data` must be a dict.") # make sure all elements in metadata are little endian meta_data = switch_endianness(meta_data, "<") - if self.ensure_channel_first: - img_array = EnsureChannelFirst()(img_array, meta_data) - if self.image_only: - return img_array meta_data[Key.FILENAME_OR_OBJ] = f"{ensure_tuple(filename)[0]}" # Path obj should be strings for data loader - - return img_array, meta_data + img = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data, self.simple_keys) + if self.ensure_channel_first: + img = EnsureChannelFirst()(img) + if self.image_only: + return img + return img, img.meta # for compatibility purpose class SaveImage(Transform): @@ -319,8 +334,8 @@ def __init__( output_ext: str = ".nii.gz", output_dtype: DtypeLike = np.float32, resample: bool = True, - mode: Union[GridSampleMode, InterpolateMode, str] = "nearest", - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = "nearest", + padding_mode: str = GridSamplePadMode.BORDER, scale: Optional[int] = None, dtype: DtypeLike = np.float64, squeeze_end_dims: bool = True, @@ -328,7 +343,7 @@ def __init__( separate_folder: bool = True, print_log: bool = True, output_format: str = "", - writer: Union[image_writer.ImageWriter, str, None] = None, + writer: Union[Type[image_writer.ImageWriter], str, None] = None, channel_dim: Optional[int] = 0, ) -> None: self.folder_layout = FolderLayout( @@ -347,7 +362,7 @@ def __init__( writer_ = locate(f"{writer}") # search dotted path if writer_ is None: raise ValueError(f"writer {writer} not found") - writer = writer_ # type: ignore + writer = writer_ self.writers = image_writer.resolve_writer(self.output_ext) if writer is None else (writer,) self.writer_obj = None @@ -389,6 +404,7 @@ def __call__(self, img: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dic img: target data content that save into file. The image should be channel-first, shape: `[C,H,W,[D]]`. meta_data: key-value pairs of metadata corresponding to the data. """ + meta_data = img.meta if isinstance(img, MetaTensor) else meta_data subject = meta_data[Key.FILENAME_OR_OBJ] if meta_data else str(self._data_index) patch_index = meta_data.get(Key.PATCH_INDEX, None) if meta_data else None filename = self.folder_layout.filename(subject=f"{subject}", idx=patch_index) diff --git a/monai/transforms/io/dictionary.py b/monai/transforms/io/dictionary.py index 30dedc7810..42918f5e63 100644 --- a/monai/transforms/io/dictionary.py +++ b/monai/transforms/io/dictionary.py @@ -16,7 +16,7 @@ """ from pathlib import Path -from typing import Optional, Union +from typing import Optional, Type, Union import numpy as np @@ -25,7 +25,7 @@ from monai.data.image_reader import ImageReader from monai.transforms.io.array import LoadImage, SaveImage from monai.transforms.transform import MapTransform -from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, ensure_tuple, ensure_tuple_rep +from monai.utils import GridSamplePadMode, ensure_tuple, ensure_tuple_rep from monai.utils.enums import PostFix __all__ = ["LoadImaged", "LoadImageD", "LoadImageDict", "SaveImaged", "SaveImageD", "SaveImageDict"] @@ -38,7 +38,7 @@ class LoadImaged(MapTransform): Dictionary-based wrapper of :py:class:`monai.transforms.LoadImage`, It can load both image data and metadata. When loading a list of files in one key, the arrays will be stacked and a new dimension will be added as the first dimension - In this case, the meta data of the first image will be used to represent the stacked result. + In this case, the metadata of the first image will be used to represent the stacked result. The affine transform of all the stacked images should be same. The output metadata field will be created as ``meta_keys`` or ``key_{meta_key_postfix}``. @@ -74,6 +74,7 @@ def __init__( overwriting: bool = False, image_only: bool = False, ensure_channel_first: bool = False, + simple_keys: bool = False, allow_missing_keys: bool = False, *args, **kwargs, @@ -82,7 +83,7 @@ def __init__( Args: keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` - reader: reader to load image file and meta data + reader: reader to load image file and metadata - if `reader` is None, a default set of `SUPPORTED_READERS` will be used. - if `reader` is a string, it's treated as a class name or dotted path (such as ``"monai.data.ITKReader"``), the supported built-in reader classes are @@ -90,25 +91,26 @@ def __init__( a reader instance will be constructed with the `*args` and `**kwargs` parameters. - if `reader` is a reader class/instance, it will be registered to this loader accordingly. dtype: if not None, convert the loaded image data to this data type. - meta_keys: explicitly indicate the key to store the corresponding meta data dictionary. - the meta data is a dictionary object which contains: filename, original_shape, etc. + meta_keys: explicitly indicate the key to store the corresponding metadata dictionary. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. meta_key_postfix: if meta_keys is None, use `key_{postfix}` to store the metadata of the nifti image, - default is `meta_dict`. The meta data is a dictionary object. + default is `meta_dict`. The metadata is a dictionary object. For example, load nifti file for `image`, store the metadata into `image_meta_dict`. - overwriting: whether allow overwriting existing meta data of same key. + overwriting: whether allow overwriting existing metadata of same key. default is False, which will raise exception if encountering existing key. image_only: if True return dictionary containing just only the image volumes, otherwise return dictionary containing image data array and header dict per input key. - ensure_channel_first: if `True` and loaded both image array and meta data, automatically convert + ensure_channel_first: if `True` and loaded both image array and metadata, automatically convert the image array shape to `channel first`. default to `False`. + simple_keys: whether to remove redundant metadata keys, default to False for backward compatibility. allow_missing_keys: don't raise exception if key is missing. args: additional parameters for reader if providing a reader name. kwargs: additional parameters for reader if providing a reader name. """ super().__init__(keys, allow_missing_keys) - self._loader = LoadImage(reader, image_only, dtype, ensure_channel_first, *args, **kwargs) + self._loader = LoadImage(reader, image_only, dtype, ensure_channel_first, simple_keys, *args, **kwargs) if not isinstance(meta_key_postfix, str): raise TypeError(f"meta_key_postfix must be a str but is {type(meta_key_postfix).__name__}.") self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) @@ -130,8 +132,6 @@ def __call__(self, data, reader: Optional[ImageReader] = None): for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): data = self._loader(d[key], reader) if self._loader.image_only: - if not isinstance(data, np.ndarray): - raise ValueError("loader must return a numpy array (because image_only=True was used).") d[key] = data else: if not isinstance(data, (tuple, list)): @@ -141,7 +141,7 @@ def __call__(self, data, reader: Optional[ImageReader] = None): raise ValueError("metadata must be a dict.") meta_key = meta_key or f"{key}_{meta_key_postfix}" if meta_key in d and not self.overwriting: - raise KeyError(f"Meta data with key {meta_key} already exists and overwriting=False.") + raise KeyError(f"Metadata with key {meta_key} already exists and overwriting=False.") d[meta_key] = data[1] return d @@ -226,8 +226,8 @@ def __init__( output_postfix: str = "trans", output_ext: str = ".nii.gz", resample: bool = True, - mode: Union[GridSampleMode, InterpolateMode, str] = "nearest", - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = "nearest", + padding_mode: str = GridSamplePadMode.BORDER, scale: Optional[int] = None, dtype: DtypeLike = np.float64, output_dtype: DtypeLike = np.float32, @@ -237,7 +237,7 @@ def __init__( separate_folder: bool = True, print_log: bool = True, output_format: str = "", - writer: Union[image_writer.ImageWriter, str, None] = None, + writer: Union[Type[image_writer.ImageWriter], str, None] = None, ) -> None: super().__init__(keys, allow_missing_keys) self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys)) @@ -268,7 +268,7 @@ def __call__(self, data): for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): if meta_key is None and meta_key_postfix is not None: meta_key = f"{key}_{meta_key_postfix}" - meta_data = d[meta_key] if meta_key is not None else None + meta_data = d.get(meta_key) if meta_key is not None else None self.saver(img=d[key], meta_data=meta_data) return d diff --git a/monai/transforms/meta_utility/__init__.py b/monai/transforms/meta_utility/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/transforms/meta_utility/__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/monai/transforms/meta_utility/dictionary.py b/monai/transforms/meta_utility/dictionary.py new file mode 100644 index 0000000000..1dcdf3483c --- /dev/null +++ b/monai/transforms/meta_utility/dictionary.py @@ -0,0 +1,106 @@ +# 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. +""" +A collection of dictionary-based wrappers for moving between MetaTensor types and dictionaries of data. +These can be used to make backwards compatible code. + +Class names are ended with 'd' to denote dictionary-based transforms. +""" + +from copy import deepcopy +from typing import Dict, Hashable, Mapping + +from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_tensor import MetaTensor +from monai.transforms.inverse import InvertibleTransform +from monai.transforms.transform import MapTransform +from monai.utils.enums import PostFix, TransformBackends + +__all__ = [ + "FromMetaTensord", + "FromMetaTensorD", + "FromMetaTensorDict", + "ToMetaTensord", + "ToMetaTensorD", + "ToMetaTensorDict", +] + + +class FromMetaTensord(MapTransform, InvertibleTransform): + """ + Dictionary-based transform to convert MetaTensor to a dictionary. + + If input is `{"a": MetaTensor, "b": MetaTensor}`, then output will + have the form `{"a": torch.Tensor, "a_meta_dict": dict, "a_transforms": list, "b": ...}`. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + im: MetaTensor = d[key] # type: ignore + d.update(im.as_dict(key)) + self.push_transform(d, key) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + for key in self.key_iterator(d): + # check transform + _ = self.get_most_recent_transform(d, key) + # do the inverse + im = d[key] + meta = d.pop(PostFix.meta(key), None) + transforms = d.pop(PostFix.transforms(key), None) + im = MetaTensor(im, meta=meta, applied_operations=transforms) # type: ignore + d[key] = im + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class ToMetaTensord(MapTransform, InvertibleTransform): + """ + Dictionary-based transform to convert a dictionary to MetaTensor. + + If input is `{"a": torch.Tensor, "a_meta_dict": dict, "b": ...}`, then output will + have the form `{"a": MetaTensor, "b": MetaTensor}`. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + self.push_transform(d, key) + im = d[key] + meta = d.pop(PostFix.meta(key), None) + transforms = d.pop(PostFix.transforms(key), None) + im = MetaTensor(im, meta=meta, applied_operations=transforms) # type: ignore + d[key] = im + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + for key in self.key_iterator(d): + # check transform + _ = self.get_most_recent_transform(d, key) + # do the inverse + im: MetaTensor = d[key] # type: ignore + d.update(im.as_dict(key)) + # Remove the applied transform + self.pop_transform(d, key) + return d + + +FromMetaTensorD = FromMetaTensorDict = FromMetaTensord +ToMetaTensorD = ToMetaTensorDict = ToMetaTensord diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 6396435aa7..29aa39d7ac 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -20,12 +20,20 @@ import torch from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_obj import get_track_meta from monai.networks import one_hot from monai.networks.layers import GaussianFilter, apply_filter from monai.transforms.transform import Transform from monai.transforms.utils import fill_holes, get_largest_connected_component_mask, get_unique_labels from monai.transforms.utils_pytorch_numpy_unification import unravel_index -from monai.utils import TransformBackends, convert_data_type, deprecated_arg, ensure_tuple, look_up_option +from monai.utils import ( + TransformBackends, + convert_data_type, + convert_to_tensor, + deprecated_arg, + ensure_tuple, + look_up_option, +) from monai.utils.type_conversion import convert_to_dst_type __all__ = [ @@ -95,6 +103,7 @@ def __call__( raise TypeError(f"other must be None or callable but is {type(other).__name__}.") # convert to float as activation must operate on float tensor + img = convert_to_tensor(img, track_meta=get_track_meta()) img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float) if sigmoid or self.sigmoid: img_t = torch.sigmoid(img_t) @@ -230,7 +239,7 @@ def __call__( if isinstance(threshold, bool): warnings.warn("`threshold_values=True/False` is deprecated, please use `threshold=value` instead.") threshold = logit_thresh if threshold else None - + img = convert_to_tensor(img, track_meta=get_track_meta()) img_t, *_ = convert_data_type(img, torch.Tensor) if argmax or self.argmax: img_t = torch.argmax(img_t, dim=0, keepdim=True) @@ -344,28 +353,29 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: applied_labels = self.applied_labels else: applied_labels = tuple(get_unique_labels(img, is_onehot, discard=0)) - + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_: torch.Tensor = convert_to_tensor(img, track_meta=False) if self.independent: for i in applied_labels: - foreground = img[i] > 0 if is_onehot else img[0] == i + foreground = img_[i] > 0 if is_onehot else img_[0] == i mask = get_largest_connected_component_mask(foreground, self.connectivity) if is_onehot: - img[i][foreground != mask] = 0 + img_[i][foreground != mask] = 0 else: - img[0][foreground != mask] = 0 - return img + img_[0][foreground != mask] = 0 + return convert_to_dst_type(img_, dst=img)[0] if not is_onehot: # not one-hot, union of labels - labels, *_ = convert_to_dst_type(applied_labels, dst=img, wrap_sequence=True) - foreground = (img[..., None] == labels).any(-1)[0] + labels, *_ = convert_to_dst_type(applied_labels, dst=img_, wrap_sequence=True) + foreground = (img_[..., None] == labels).any(-1)[0] mask = get_largest_connected_component_mask(foreground, self.connectivity) - img[0][foreground != mask] = 0 - return img + img_[0][foreground != mask] = 0 + return convert_to_dst_type(img_, dst=img)[0] # one-hot, union of labels - foreground = (img[applied_labels, ...] == 1).any(0) + foreground = (img_[applied_labels, ...] == 1).any(0) mask = get_largest_connected_component_mask(foreground, self.connectivity) for i in applied_labels: - img[i][foreground != mask] = 0 - return img + img_[i][foreground != mask] = 0 + return convert_to_dst_type(img_, dst=img)[0] class LabelFilter: @@ -414,13 +424,15 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: raise NotImplementedError(f"{self.__class__} can not handle data of type {type(img)}.") if isinstance(img, torch.Tensor): + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_ = convert_to_tensor(img, track_meta=False) if hasattr(torch, "isin"): # `isin` is new in torch 1.10.0 - appl_lbls = torch.as_tensor(self.applied_labels, device=img.device) - return torch.where(torch.isin(img, appl_lbls), img, torch.tensor(0.0).to(img)) - else: - out = self(img.detach().cpu().numpy()) - out, *_ = convert_to_dst_type(out, img) - return out + appl_lbls = torch.as_tensor(self.applied_labels, device=img_.device) + out = torch.where(torch.isin(img_, appl_lbls), img_, torch.tensor(0.0).to(img_)) + return convert_to_dst_type(out, dst=img)[0] + out: NdarrayOrTensor = self(img_.detach().cpu().numpy()) # type: ignore + out = convert_to_dst_type(out, img)[0] # type: ignore + return out return np.asarray(np.where(np.isin(img, self.applied_labels), img, 0)) @@ -497,8 +509,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Returns: Pytorch Tensor or numpy array of shape [C, spatial_dim1[, spatial_dim2, ...]]. """ - if not isinstance(img, (np.ndarray, torch.Tensor)): - raise NotImplementedError(f"{self.__class__} can not handle data of type {type(img)}.") + img = convert_to_tensor(img, track_meta=get_track_meta()) img_np, *_ = convert_data_type(img, np.ndarray) out_np: np.ndarray = fill_holes(img_np, self.applied_labels, self.connectivity) out, *_ = convert_to_dst_type(out_np, img) @@ -541,7 +552,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: ideally the edge should be thin enough, but now it has a thickness. """ - img_: torch.Tensor = convert_data_type(img, torch.Tensor)[0] + 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 img_ = img_.unsqueeze(0) # adds a batch dim if spatial_dims == 2: @@ -733,7 +745,7 @@ def __call__(self, prob_map: NdarrayOrTensor): if self.sigma != 0: if not isinstance(prob_map, torch.Tensor): prob_map = torch.as_tensor(prob_map, dtype=torch.float) - self.filter.to(prob_map) + self.filter.to(prob_map.device) prob_map = self.filter(prob_map) prob_map_shape = prob_map.shape diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index 00ffe7edf7..3704d92ec3 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -23,6 +23,7 @@ from monai.config.type_definitions import KeysCollection, NdarrayOrTensor, PathLike from monai.data.csv_saver import CSVSaver +from monai.data.meta_tensor import MetaTensor from monai.transforms.inverse import InvertibleTransform from monai.transforms.post.array import ( Activations, @@ -37,8 +38,8 @@ ) from monai.transforms.transform import MapTransform from monai.transforms.utility.array import ToTensor -from monai.transforms.utils import allow_missing_keys_mode, convert_inverse_interp_mode -from monai.utils import deprecated_arg, ensure_tuple, ensure_tuple_rep +from monai.transforms.utils import allow_missing_keys_mode, convert_applied_interp_mode +from monai.utils import convert_to_tensor, deprecated_arg, ensure_tuple, ensure_tuple_rep from monai.utils.enums import PostFix __all__ = [ @@ -160,7 +161,7 @@ def __init__( it also can be a sequence of bool, each element corresponds to a key in ``keys``. to_onehot: if not None, convert input data into the one-hot format with specified number of classes. defaults to ``None``. it also can be a sequence, each element corresponds to a key in ``keys``. - threshold: if not None, threshold the float values to int number 0 or 1 with specified theashold value. + threshold: if not None, threshold the float values to int number 0 or 1 with specified threshold value. defaults to ``None``. it also can be a sequence, each element corresponds to a key in ``keys``. rounding: if not None, round the data according to the specified option, available options: ["torchrounding"]. it also can be a sequence of str or None, @@ -543,7 +544,7 @@ def __init__( self, keys: KeysCollection, transform: InvertibleTransform, - orig_keys: KeysCollection, + orig_keys: Optional[KeysCollection] = None, meta_keys: Optional[KeysCollection] = None, orig_meta_keys: Optional[KeysCollection] = None, meta_key_postfix: str = DEFAULT_POST_FIX, @@ -558,20 +559,20 @@ def __init__( keys: the key of expected data in the dict, the inverse of ``transforms`` will be applied on it in-place. It also can be a list of keys, will apply the inverse transform respectively. transform: the transform applied to ``orig_key``, its inverse will be applied on ``key``. - orig_keys: the key of the original input data in the dict. + orig_keys: the key of the original input data in the dict. These keys default to `self.keys` if not set. the transform trace information of ``transforms`` should be stored at ``{orig_keys}_transforms``. It can also be a list of keys, each matches the ``keys``. - meta_keys: The key to output the inverted meta data dictionary. - The meta data is a dictionary optionally containing: filename, original_shape. + meta_keys: The key to output the inverted metadata dictionary. + The metadata is a dictionary optionally containing: filename, original_shape. It can be a sequence of strings, maps to ``keys``. - If None, will try to create a meta data dict with the default key: `{key}_{meta_key_postfix}`. - orig_meta_keys: the key of the meta data of original input data. - The meta data is a dictionary optionally containing: filename, original_shape. + If None, will try to create a metadata dict with the default key: `{key}_{meta_key_postfix}`. + orig_meta_keys: the key of the metadata of original input data. + The metadata is a dictionary optionally containing: filename, original_shape. It can be a sequence of strings, maps to the `keys`. - If None, will try to create a meta data dict with the default key: `{orig_key}_{meta_key_postfix}`. - This meta data dict will also be included in the inverted dict, stored in `meta_keys`. + If None, will try to create a metadata dict with the default key: `{orig_key}_{meta_key_postfix}`. + This metadata dict will also be included in the inverted dict, stored in `meta_keys`. meta_key_postfix: if `orig_meta_keys` is None, use `{orig_key}_{meta_key_postfix}` to fetch the - meta data from dict, if `meta_keys` is None, use `{key}_{meta_key_postfix}`. Default: ``"meta_dict"``. + metadata from dict, if `meta_keys` is None, use `{key}_{meta_key_postfix}`. Default: ``"meta_dict"``. nearest_interp: whether to use `nearest` interpolation mode when inverting the spatial transforms, default to `True`. If `False`, use the same interpolation mode as the original transform. It also can be a list of bool, each matches to the `keys` data. @@ -588,7 +589,7 @@ def __init__( if not isinstance(transform, InvertibleTransform): raise ValueError("transform is not invertible, can't invert transform for the data.") self.transform = transform - self.orig_keys = ensure_tuple_rep(orig_keys, len(self.keys)) + self.orig_keys = ensure_tuple_rep(orig_keys, len(self.keys)) if orig_keys is not None else self.keys self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) if len(self.keys) != len(self.meta_keys): raise ValueError("meta_keys should have the same length as keys.") @@ -623,44 +624,59 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]: self.device, self.post_func, ): - transform_key = InvertibleTransform.trace_key(orig_key) - if transform_key not in d: - warnings.warn(f"transform info of `{orig_key}` is not available or no InvertibleTransform applied.") - continue - - transform_info = d[transform_key] + if isinstance(d[key], MetaTensor): + if orig_key not in d: + warnings.warn(f"transform info of `{orig_key}` is not available in MetaTensor {key}.") + continue + else: + transform_key = InvertibleTransform.trace_key(orig_key) + if transform_key not in d: + warnings.warn(f"transform info of `{orig_key}` is not available or no InvertibleTransform applied.") + continue + + if orig_key in d and isinstance(d[orig_key], MetaTensor): + transform_info = d[orig_key].applied_operations + meta_info = d[orig_key].meta + else: + transform_info = d[InvertibleTransform.trace_key(orig_key)] + meta_info = d.get(orig_meta_key or f"{orig_key}_{meta_key_postfix}", {}) if nearest_interp: - transform_info = convert_inverse_interp_mode( + transform_info = convert_applied_interp_mode( trans_info=deepcopy(transform_info), mode="nearest", align_corners=None ) - input = d[key] - if isinstance(input, torch.Tensor): - input = input.detach() + inputs = d[key] + if isinstance(inputs, torch.Tensor): + inputs = inputs.detach() + + if not isinstance(inputs, MetaTensor): + inputs = convert_to_tensor(inputs, track_meta=True) + inputs.applied_operations = transform_info + inputs.meta = meta_info # construct the input dict data - input_dict = {orig_key: input, transform_key: transform_info} - orig_meta_key = orig_meta_key or f"{orig_key}_{meta_key_postfix}" - if orig_meta_key in d: - input_dict[orig_meta_key] = d[orig_meta_key] + input_dict = {orig_key: inputs} with allow_missing_keys_mode(self.transform): # type: ignore inverted = self.transform.inverse(input_dict) # save the inverted data - d[key] = post_func(self._totensor(inverted[orig_key]).to(device) if to_tensor else inverted[orig_key]) + if to_tensor and not isinstance(inverted[orig_key], MetaTensor): + inverted_data = self._totensor(inverted[orig_key]) + else: + inverted_data = inverted[orig_key] + d[key] = post_func(inverted_data.to(device)) # save the inverted meta dict if orig_meta_key in d: meta_key = meta_key or f"{key}_{meta_key_postfix}" d[meta_key] = inverted.get(orig_meta_key) - return d class SaveClassificationd(MapTransform): """ - Save the classification results and meta data into CSV file or other storage. + Save the classification results and metadata into CSV file or other storage. """ @@ -681,16 +697,16 @@ def __init__( Args: keys: keys of the corresponding items to model output, this transform only supports 1 key. See also: :py:class:`monai.transforms.compose.MapTransform` - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. will extract the filename of input image to save classification results. meta_key_postfix: `key_{postfix}` was used to store the metadata in `LoadImaged`. so need the key to extract the metadata of input image, like filename, etc. default is `meta_dict`. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. this arg only works when `meta_keys=None`. if no corresponding metadata, set to `None`. saver: the saver instance to save classification results, if None, create a CSVSaver internally. the saver must provide `save(data, meta_data)` and `finalize()` APIs. diff --git a/monai/transforms/smooth_field/array.py b/monai/transforms/smooth_field/array.py index f581687ea5..953c589288 100644 --- a/monai/transforms/smooth_field/array.py +++ b/monai/transforms/smooth_field/array.py @@ -17,8 +17,8 @@ import torch from torch.nn.functional import grid_sample, interpolate -import monai from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_obj import get_track_meta from monai.networks.utils import meshgrid_ij from monai.transforms.transform import Randomizable, RandomizableTransform from monai.transforms.utils_pytorch_numpy_unification import moveaxis @@ -61,7 +61,7 @@ def __init__( high: float = 1.0, channels: int = 1, spatial_size: Optional[Sequence[int]] = None, - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, + mode: str = InterpolateMode.AREA, align_corners: Optional[bool] = None, device: Optional[torch.device] = None, ): @@ -109,7 +109,7 @@ def set_spatial_size(self, spatial_size: Optional[Sequence[int]]) -> None: self.spatial_size = tuple(spatial_size) self.spatial_zoom = tuple(s / f for s, f in zip(self.spatial_size, self.total_rand_size)) - def set_mode(self, mode: Union[monai.utils.InterpolateMode, str]) -> None: + def set_mode(self, mode: str) -> None: self.mode = mode def __call__(self, randomize=False) -> torch.Tensor: @@ -119,10 +119,10 @@ def __call__(self, randomize=False) -> torch.Tensor: field = self.field.clone() if self.spatial_zoom is not None: - resized_field = interpolate( # type: ignore + resized_field = interpolate( input=field, scale_factor=self.spatial_zoom, - mode=look_up_option(self.mode, InterpolateMode).value, + mode=look_up_option(self.mode, InterpolateMode), align_corners=self.align_corners, recompute_scale_factor=False, ) @@ -147,7 +147,7 @@ class RandSmoothFieldAdjustContrast(RandomizableTransform): edges of the input volume of that width will be mostly unchanged. Contrast is changed by raising input values by the power of the smooth field so the range of values given by `gamma` should be chosen with this in mind. For example, a minimum value of 0 in `gamma` will produce white areas so this should be avoided. - Afte the contrast is adjusted the values of the result are rescaled to the range of the original input. + After the contrast is adjusted the values of the result are rescaled to the range of the original input. Args: spatial_size: size of input array's spatial dimensions @@ -167,7 +167,7 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, + mode: str = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, gamma: Union[Sequence[float], float] = (0.5, 4.5), @@ -209,13 +209,14 @@ def randomize(self, data: Optional[Any] = None) -> None: if self._do_transform: self.sfield.randomize() - def set_mode(self, mode: Union[monai.utils.InterpolateMode, str]) -> None: + def set_mode(self, mode: str) -> None: self.sfield.set_mode(mode) def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: """ Apply the transform to `img`, if `randomize` randomizing the smooth field otherwise reusing the previous. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -267,7 +268,7 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, + mode: str = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, gamma: Union[Sequence[float], float] = (0.1, 1.0), @@ -309,13 +310,14 @@ def randomize(self, data: Optional[Any] = None) -> None: if self._do_transform: self.sfield.randomize() - def set_mode(self, mode: Union[InterpolateMode, str]) -> None: + def set_mode(self, mode: str) -> None: self.sfield.set_mode(mode) def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: """ Apply the transform to `img`, if `randomize` randomizing the smooth field otherwise reusing the previous. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -363,13 +365,13 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - field_mode: Union[InterpolateMode, str] = InterpolateMode.AREA, + field_mode: str = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, def_range: Union[Sequence[float], float] = 1.0, grid_dtype=torch.float32, - grid_mode: Union[GridSampleMode, str] = GridSampleMode.NEAREST, - grid_padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + grid_mode: str = GridSampleMode.NEAREST, + grid_padding_mode: str = GridSamplePadMode.BORDER, grid_align_corners: Optional[bool] = False, device: Optional[torch.device] = None, ): @@ -422,15 +424,16 @@ def randomize(self, data: Optional[Any] = None) -> None: if self._do_transform: self.sfield.randomize() - def set_field_mode(self, mode: Union[monai.utils.InterpolateMode, str]) -> None: + def set_field_mode(self, mode: str) -> None: self.sfield.set_mode(mode) - def set_grid_mode(self, mode: Union[monai.utils.GridSampleMode, str]) -> None: + def set_grid_mode(self, mode: str) -> None: self.grid_mode = mode def __call__( self, img: NdarrayOrTensor, randomize: bool = True, device: Optional[torch.device] = None ) -> NdarrayOrTensor: + img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() @@ -449,9 +452,9 @@ def __call__( out = grid_sample( input=img_t, grid=dgrid, - mode=look_up_option(self.grid_mode, GridSampleMode).value, + mode=look_up_option(self.grid_mode, GridSampleMode), align_corners=self.grid_align_corners, - padding_mode=look_up_option(self.grid_padding_mode, GridSamplePadMode).value, + padding_mode=look_up_option(self.grid_padding_mode, GridSamplePadMode), ) out_t, *_ = convert_to_dst_type(out.squeeze(0), img) diff --git a/monai/transforms/smooth_field/dictionary.py b/monai/transforms/smooth_field/dictionary.py index 24890140cc..48e00b9e4a 100644 --- a/monai/transforms/smooth_field/dictionary.py +++ b/monai/transforms/smooth_field/dictionary.py @@ -15,15 +15,16 @@ import numpy as np import torch -from monai.config import KeysCollection +from monai.config import KeysCollection, SequenceStr from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_obj import get_track_meta from monai.transforms.smooth_field.array import ( RandSmoothDeform, RandSmoothFieldAdjustContrast, RandSmoothFieldAdjustIntensity, ) from monai.transforms.transform import MapTransform, RandomizableTransform -from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, ensure_tuple_rep +from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, convert_to_tensor, ensure_tuple_rep from monai.utils.enums import TransformBackends __all__ = [ @@ -39,10 +40,6 @@ ] -InterpolateModeType = Union[InterpolateMode, str] -GridSampleModeType = Union[GridSampleMode, str] - - class RandSmoothFieldAdjustContrastd(RandomizableTransform, MapTransform): """ Dictionary version of RandSmoothFieldAdjustContrast. @@ -71,7 +68,7 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - mode: Union[InterpolateModeType, Sequence[InterpolateModeType]] = InterpolateMode.AREA, + mode: SequenceStr = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, gamma: Union[Sequence[float], float] = (0.5, 4.5), @@ -108,11 +105,11 @@ def randomize(self, data: Optional[Any] = None) -> None: def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]: self.randomize() - - if not self._do_transform: - return data - d = dict(data) + if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + return d for idx, key in enumerate(self.key_iterator(d)): self.trans.set_mode(self.mode[idx % len(self.mode)]) @@ -149,7 +146,7 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - mode: Union[InterpolateModeType, Sequence[InterpolateModeType]] = InterpolateMode.AREA, + mode: SequenceStr = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, gamma: Union[Sequence[float], float] = (0.1, 1.0), @@ -185,10 +182,11 @@ def randomize(self, data: Optional[Any] = None) -> None: def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]: self.randomize() - if not self._do_transform: - return data - d = dict(data) + if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + return d for idx, key in enumerate(self.key_iterator(d)): self.trans.set_mode(self.mode[idx % len(self.mode)]) @@ -229,13 +227,13 @@ def __init__( spatial_size: Sequence[int], rand_size: Sequence[int], pad: int = 0, - field_mode: Union[InterpolateModeType, Sequence[InterpolateModeType]] = InterpolateMode.AREA, + field_mode: SequenceStr = InterpolateMode.AREA, align_corners: Optional[bool] = None, prob: float = 0.1, def_range: Union[Sequence[float], float] = 1.0, grid_dtype=torch.float32, - grid_mode: Union[GridSampleModeType, Sequence[GridSampleModeType]] = GridSampleMode.NEAREST, - grid_padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + grid_mode: SequenceStr = GridSampleMode.NEAREST, + grid_padding_mode: str = GridSamplePadMode.BORDER, grid_align_corners: Optional[bool] = False, device: Optional[torch.device] = None, ): @@ -274,10 +272,11 @@ def randomize(self, data: Optional[Any] = None) -> None: def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]: self.randomize() - if not self._do_transform: - return data - d = dict(data) + if not self._do_transform: + for key in self.key_iterator(d): + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + return d for idx, key in enumerate(self.key_iterator(d)): self.trans.set_field_mode(self.field_mode[idx % len(self.field_mode)]) diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py index 41d048fc19..d875d3adee 100644 --- a/monai/transforms/spatial/array.py +++ b/monai/transforms/spatial/array.py @@ -14,19 +14,26 @@ """ import warnings from copy import deepcopy -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union +from enum import Enum +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import numpy as np import torch +from numpy.lib.stride_tricks import as_strided from monai.config import USE_COMPILED, DtypeLike from monai.config.type_definitions import NdarrayOrTensor -from monai.data.utils import AFFINE_TOL, compute_shape_offset, reorient_spatial_axes, to_affine_nd, zoom_affine +from monai.data.meta_obj import get_track_meta +from monai.data.meta_tensor import MetaTensor +from monai.data.utils import AFFINE_TOL, 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, normalize_transform -from monai.transforms.croppad.array import CenterSpatialCrop, Pad -from monai.transforms.transform import Randomizable, RandomizableTransform, ThreadUnsafe, Transform +from monai.transforms.croppad.array import CenterSpatialCrop, ResizeWithPadOrCrop +from monai.transforms.intensity.array import GaussianSmooth +from monai.transforms.inverse import InvertibleTransform +from monai.transforms.transform import Randomizable, RandomizableTransform, Transform from monai.transforms.utils import ( + convert_pad_mode, create_control_grid, create_grid, create_rotate, @@ -34,14 +41,16 @@ create_shear, create_translate, map_spatial_axes, + scale_affine, ) -from monai.transforms.utils_pytorch_numpy_unification import allclose, moveaxis +from monai.transforms.utils_pytorch_numpy_unification import allclose, linalg_inv, moveaxis from monai.utils import ( GridSampleMode, GridSamplePadMode, InterpolateMode, NumpyPadMode, - PytorchPadMode, + convert_to_dst_type, + convert_to_tensor, ensure_tuple, ensure_tuple_rep, ensure_tuple_size, @@ -51,9 +60,10 @@ pytorch_after, ) from monai.utils.deprecate_utils import deprecated_arg -from monai.utils.enums import TransformBackends +from monai.utils.enums import GridPatchSort, PytorchPadMode, TraceKeys, TransformBackends +from monai.utils.misc import ImageMetaKey as Key from monai.utils.module import look_up_option -from monai.utils.type_conversion import convert_data_type, convert_to_dst_type +from monai.utils.type_conversion import convert_data_type, get_equivalent_dtype, get_torch_dtype_from_string nib, has_nib = optional_import("nibabel") @@ -64,6 +74,9 @@ "Orientation", "Flip", "GridDistortion", + "GridSplit", + "GridPatch", + "RandGridPatch", "Resize", "Rotate", "Zoom", @@ -87,7 +100,7 @@ RandRange = Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] -class SpatialResample(Transform): +class SpatialResample(InvertibleTransform): """ Resample input image from the orientation/spacing defined by ``src_affine`` affine matrix into the ones specified by ``dst_affine`` affine matrix. @@ -100,8 +113,8 @@ class SpatialResample(Transform): def __init__( self, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: DtypeLike = np.float64, ): @@ -127,24 +140,62 @@ def __init__( self.align_corners = align_corners self.dtype = dtype + def _post_process( + self, + img: torch.Tensor, + src_affine: torch.Tensor, + dst_affine: torch.Tensor, + mode, + padding_mode, + align_corners, + original_spatial_shape, + ) -> torch.Tensor: + """ + Small fn to simplify returning data. If `MetaTensor`, update affine. Elif + tracking metadata is desired, create `MetaTensor` with affine. Else, return + image as `torch.Tensor`. Output type is always `torch.float32`. + + Also append the transform to the stack. + """ + dtype = img.dtype + img = convert_to_tensor(img, track_meta=get_track_meta(), dtype=torch.float32) + if get_track_meta(): + self.update_meta(img, dst_affine) + self.push_transform( + img, + extra_info={ + "dtype": str(dtype)[6:], # dtype as string; remove "torch": torch.float32 -> float32 + "mode": mode.value if isinstance(mode, Enum) else mode, + "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, + "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, + "src_affine": src_affine, + }, + orig_size=original_spatial_shape, + ) + return img + + def update_meta(self, img, dst_affine): + img.affine = dst_affine + + @deprecated_arg( + name="src_affine", since="0.9", msg_suffix="img should be `MetaTensor`, so affine can be extracted directly." + ) def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, src_affine: Optional[NdarrayOrTensor] = None, - dst_affine: Optional[NdarrayOrTensor] = None, - spatial_size: Optional[Union[Sequence[int], np.ndarray, int]] = None, - mode: Union[GridSampleMode, str, None] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str, None] = GridSamplePadMode.BORDER, + dst_affine: Optional[torch.Tensor] = None, + spatial_size: Optional[Union[Sequence[int], torch.Tensor, int]] = None, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = False, dtype: DtypeLike = None, - ) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + ) -> torch.Tensor: """ Args: img: input image to be resampled. It currently supports channel-first arrays with at most three spatial dimensions. - src_affine: source affine matrix. Defaults to ``None``, which means the identity matrix. - the shape should be `(r+1, r+1)` where `r` is the spatial rank of ``img``. - dst_affine: destination affine matrix. Defaults to ``None``, which means the same as `src_affine`. + dst_affine: destination affine matrix. Defaults to ``None``, which means the same as `img.affine`. the shape should be `(r+1, r+1)` where `r` is the spatial rank of ``img``. when `dst_affine` and `spatial_size` are None, the input will be returned without resampling, but the data type will be `float32`. @@ -173,85 +224,77 @@ def __call__( MONAI's resampling implementation will be used. Set `dst_affine` and `spatial_size` to `None` to turn off the resampling step. """ - if src_affine is None: - src_affine = np.eye(4, dtype=np.float64) - spatial_rank = min(len(img.shape) - 1, src_affine.shape[0] - 1, 3) + # get dtype as torch (e.g., torch.float64) + _dtype = get_equivalent_dtype(dtype or self.dtype or img.dtype, torch.Tensor) + align_corners = self.align_corners if align_corners is None else align_corners + mode = mode or self.mode + padding_mode = padding_mode or self.padding_mode + original_spatial_shape = img.shape[1:] + + src_affine_: torch.Tensor = img.affine if isinstance(img, MetaTensor) else torch.eye(4) + img = convert_to_tensor(data=img, track_meta=get_track_meta(), dtype=_dtype) + spatial_rank = min(len(img.shape) - 1, src_affine_.shape[0] - 1, 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 - src_affine = to_affine_nd(spatial_rank, src_affine) - 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, dst_affine, dtype=torch.float32) + src_affine_ = to_affine_nd(spatial_rank, src_affine_).to(_dtype) + 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] + if not isinstance(dst_affine, torch.Tensor): + raise ValueError(f"dst_affine should be a torch.Tensor, got {type(dst_affine)}") - in_spatial_size = np.asarray(img.shape[1 : spatial_rank + 1]) + in_spatial_size = torch.tensor(img.shape[1 : spatial_rank + 1]) if isinstance(spatial_size, int) and (spatial_size == -1): # using the input spatial size spatial_size = in_spatial_size elif spatial_size is None and spatial_rank > 1: # auto spatial size - spatial_size, _ = compute_shape_offset(in_spatial_size, src_affine, dst_affine) # type: ignore - spatial_size = np.asarray(fall_back_tuple(ensure_tuple(spatial_size)[:spatial_rank], in_spatial_size)) + spatial_size, _ = compute_shape_offset(in_spatial_size, src_affine_, dst_affine) # type: ignore + spatial_size = torch.tensor(fall_back_tuple(ensure_tuple(spatial_size)[:spatial_rank], in_spatial_size)) if ( - allclose(src_affine, dst_affine, atol=AFFINE_TOL) + allclose(src_affine_, dst_affine, atol=AFFINE_TOL) and allclose(spatial_size, in_spatial_size) or spatial_rank == 1 ): # no significant change, return original image - output_data, *_ = convert_to_dst_type(img, img, dtype=torch.float32) - return output_data, dst_affine - - if has_nib and isinstance(img, np.ndarray): - spatial_ornt, dst_r = reorient_spatial_axes(img.shape[1 : spatial_rank + 1], src_affine, dst_affine) - if allclose(dst_r, dst_affine, atol=AFFINE_TOL) and allclose(spatial_size, in_spatial_size): - # simple reorientation achieves the desired affine - spatial_ornt[:, 0] += 1 - spatial_ornt = np.concatenate([np.array([[0, 1]]), spatial_ornt]) - img_ = nib.orientations.apply_orientation(img, spatial_ornt) - output_data, *_ = convert_to_dst_type(img_, img, dtype=torch.float32) - return output_data, dst_affine + return self._post_process( + img, src_affine_, src_affine_, mode, padding_mode, align_corners, original_spatial_shape + ) try: - src_affine, *_ = convert_to_dst_type(src_affine, dst_affine) - if isinstance(src_affine, np.ndarray): - xform = np.linalg.solve(src_affine, dst_affine) - else: - xform = ( - torch.linalg.solve(src_affine, dst_affine) - if pytorch_after(1, 8, 0) - else torch.solve(dst_affine, src_affine).solution # type: ignore - ) + _s = convert_to_tensor(src_affine_, track_meta=False, device=torch.device("cpu")) + _d = convert_to_tensor(dst_affine, track_meta=False, device=torch.device("cpu")) + xform = ( + torch.linalg.solve(_s, _d) if pytorch_after(1, 8, 0) else torch.solve(_d, _s).solution # type: ignore + ) except (np.linalg.LinAlgError, RuntimeError) as e: - raise ValueError(f"src affine is not invertible: {src_affine}") from e - xform = to_affine_nd(spatial_rank, xform) + raise ValueError("src affine is not invertible.") from e + xform = to_affine_nd(spatial_rank, xform).to(device=img.device, dtype=_dtype) # no resampling if it's identity transform - if allclose(xform, np.diag(np.ones(len(xform))), atol=AFFINE_TOL) and allclose(spatial_size, in_spatial_size): - output_data, *_ = convert_to_dst_type(img, img, dtype=torch.float32) - return output_data, dst_affine + if allclose(xform, torch.eye(len(xform)), atol=AFFINE_TOL) and allclose(spatial_size, in_spatial_size): + return self._post_process( + img, src_affine_, src_affine_, mode, padding_mode, align_corners, original_spatial_shape + ) - _dtype = dtype or self.dtype or img.dtype - in_spatial_size = in_spatial_size.tolist() + in_spatial_size = in_spatial_size.tolist() # type: ignore chns, additional_dims = img.shape[0], img.shape[spatial_rank + 1 :] # beyond three spatial dims - # resample - img_ = convert_data_type(img, torch.Tensor, dtype=_dtype)[0] - xform = convert_to_dst_type(xform, img_)[0] - align_corners = self.align_corners if align_corners is None else align_corners - mode = mode or self.mode - padding_mode = padding_mode or self.padding_mode + if additional_dims: xform_shape = [-1] + in_spatial_size - img_ = img_.reshape(xform_shape) + img = img.reshape(xform_shape) # type: ignore if align_corners: - _t_r = torch.diag(torch.ones(len(xform), dtype=xform.dtype, device=xform.device)) # type: ignore + _t_r = torch.eye(len(xform), dtype=xform.dtype, device=xform.device) for idx, d_dst in enumerate(spatial_size[:spatial_rank]): _t_r[idx, -1] = (max(d_dst, 2) - 1.0) / 2.0 xform = xform @ _t_r if not USE_COMPILED: _t_l = normalize_transform( in_spatial_size, xform.device, xform.dtype, align_corners=True # type: ignore - ) - xform = _t_l @ xform # type: ignore + )[0] + xform = _t_l @ xform affine_xform = Affine( - affine=xform, spatial_size=spatial_size, norm_coords=False, image_only=True, dtype=_dtype + affine=xform, spatial_size=spatial_size, normalized=True, image_only=True, dtype=_dtype ) - output_data = affine_xform(img_, mode=mode, padding_mode=padding_mode) + with affine_xform.trace_transform(False): + img = affine_xform(img, mode=mode, padding_mode=padding_mode) else: affine_xform = AffineTransform( normalized=False, @@ -260,67 +303,125 @@ def __call__( align_corners=align_corners, reverse_indexing=True, ) - output_data = affine_xform(img_.unsqueeze(0), theta=xform, spatial_size=spatial_size).squeeze(0) + img = affine_xform(img.unsqueeze(0), theta=xform, spatial_size=spatial_size).squeeze(0) if additional_dims: full_shape = (chns, *spatial_size, *additional_dims) - output_data = output_data.reshape(full_shape) - # output dtype float - output_data, *_ = convert_to_dst_type(output_data, img, dtype=torch.float32) - return output_data, dst_affine + img = img.reshape(full_shape) + + return self._post_process( + img, src_affine_, dst_affine, mode, padding_mode, align_corners, original_spatial_shape + ) + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + # Create inverse transform + kw_args = transform[TraceKeys.EXTRA_INFO] + # need to convert dtype from string back to torch.dtype + kw_args["dtype"] = get_torch_dtype_from_string(kw_args["dtype"]) + # source becomes destination + kw_args["dst_affine"] = kw_args.pop("src_affine") + kw_args["spatial_size"] = transform[TraceKeys.ORIG_SIZE] + if kw_args.get("align_corners") == TraceKeys.NONE: + kw_args["align_corners"] = False + with self.trace_transform(False): + # we can't use `self.__call__` in case a child class calls this inverse. + out: torch.Tensor = SpatialResample.__call__(self, data, **kw_args) + return out class ResampleToMatch(SpatialResample): - """Resample an image to match given meta data. The affine matrix will be aligned, + """Resample an image to match given metadata. The affine matrix will be aligned, and the size of the output image will match.""" - def __call__( # type: ignore + def update_meta(self, img: torch.Tensor, dst_affine=None, img_dst=None): + if dst_affine is not None: + super().update_meta(img, dst_affine) + if isinstance(img_dst, MetaTensor) and isinstance(img, MetaTensor): + original_fname = img.meta[Key.FILENAME_OR_OBJ] + img.meta = deepcopy(img_dst.meta) + img.meta[Key.FILENAME_OR_OBJ] = original_fname # keep the original name, the others are overwritten + + @deprecated_arg( + name="src_meta", since="0.9", msg_suffix="img should be `MetaTensor`, so affine can be extracted directly." + ) + @deprecated_arg( + name="dst_meta", since="0.9", msg_suffix="img_dst should be `MetaTensor`, so affine can be extracted directly." + ) + def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, + img_dst: torch.Tensor, src_meta: Optional[Dict] = None, dst_meta: Optional[Dict] = None, - mode: Union[GridSampleMode, str, None] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str, None] = GridSamplePadMode.BORDER, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = False, dtype: DtypeLike = None, - ): - if src_meta is None: - raise RuntimeError("`in_meta` is missing") - if dst_meta is None: - raise RuntimeError("`out_meta` is missing") - mode = mode or self.mode - padding_mode = padding_mode or self.padding_mode - align_corners = self.align_corners if align_corners is None else align_corners - dtype = dtype or self.dtype - src_affine = src_meta.get("affine") - dst_affine = dst_meta.get("affine") - img, updated_affine = super().__call__( + ) -> torch.Tensor: + """ + Args: + img: input image to be resampled to match ``dst_meta``. It currently supports channel-first arrays with + at most three spatial dimensions. + src_meta: Dictionary containing the source affine matrix in the form ``{'affine':src_affine}``. + If ``affine`` is not specified, an identity matrix is assumed. Defaults to ``None``. + See also: https://docs.monai.io/en/stable/transforms.html#spatialresample + dst_meta: Dictionary containing the target affine matrix and target spatial shape in the form + ``{'affine':src_affine, 'spatial_shape':spatial_size}``. If ``affine`` is not + specified, ``src_affine`` is assumed. If ``spatial_shape`` is not specified, spatial size is + automatically computed, containing the previous field of view. Defaults to ``None``. + See also: https://docs.monai.io/en/stable/transforms.html#spatialresample + mode: {``"bilinear"``, ``"nearest"``} + Interpolation mode to calculate output values. Defaults to ``"bilinear"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html + When `USE_COMPILED` is `True`, this argument uses + ``"nearest"``, ``"bilinear"``, ``"bicubic"`` to indicate 0, 1, 3 order interpolations. + See also: https://docs.monai.io/en/stable/networks.html#grid-pull + padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``} + Padding mode for outside grid values. Defaults to ``"border"``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html + align_corners: Geometrically, we consider the pixels of the input as squares rather than points. + Defaults to ``False``. + See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html + dtype: data type for resampling computation. Defaults to ``self.dtype`` or + ``np.float64`` (for best precision). If ``None``, use the data type of input data. + To be compatible with other modules, the output data type is always `float32`. + Raises: + RuntimeError: When ``src_meta`` is missing. + RuntimeError: When ``dst_meta`` is missing. + ValueError: When the affine matrix of the source image is not invertible. + Returns: + Resampled input image, Metadata + """ + if img_dst is None: + raise RuntimeError("`img_dst` is missing.") + dst_affine = img_dst.affine if isinstance(img_dst, MetaTensor) else torch.eye(4) + img = super().__call__( img=img, - src_affine=src_affine, dst_affine=dst_affine, - spatial_size=dst_meta.get("spatial_shape"), + spatial_size=img_dst.shape[1:], # skip channel mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype, ) - dst_meta = deepcopy(dst_meta) - dst_meta["affine"] = updated_affine - return img, dst_meta + self.update_meta(img, dst_affine=dst_affine, img_dst=img_dst) + return img -class Spacing(Transform): +class Spacing(InvertibleTransform): """ Resample input image into the specified `pixdim`. """ backend = SpatialResample.backend + @deprecated_arg(name="image_only", since="0.9") def __init__( self, pixdim: Union[Sequence[float], float, np.ndarray], diagonal: bool = False, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: DtypeLike = np.float64, image_only: bool = False, @@ -359,12 +460,10 @@ def __init__( dtype: data type for resampling computation. Defaults to ``np.float64`` for best precision. If None, use the data type of input data. To be compatible with other modules, the output data type is always ``np.float32``. - image_only: return just the image or the image, the old affine and new affine. Default is `False`. """ self.pixdim = np.array(ensure_tuple(pixdim), dtype=np.float64) self.diagonal = diagonal - self.image_only = image_only self.sp_resample = SpatialResample( mode=look_up_option(mode, GridSampleMode), @@ -373,20 +472,20 @@ def __init__( dtype=dtype, ) + @deprecated_arg(name="affine", since="0.9", msg_suffix="Not needed, input should be `MetaTensor`.") def __call__( self, - data_array: NdarrayOrTensor, + data_array: torch.Tensor, affine: Optional[NdarrayOrTensor] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = None, dtype: DtypeLike = None, output_spatial_shape: Optional[Union[Sequence[int], np.ndarray, int]] = None, - ) -> Union[NdarrayOrTensor, Tuple[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor]]: + ) -> torch.Tensor: """ Args: data_array: in shape (num_channels, H[, W, ...]). - affine (matrix): (N+1)x(N+1) original affine matrix for spatially ND `data_array`. Defaults to identity. mode: {``"bilinear"``, ``"nearest"``} Interpolation mode to calculate output values. Defaults to ``self.mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html @@ -413,16 +512,22 @@ def __call__( data_array (resampled into `self.pixdim`), original affine, current affine. """ - sr = int(data_array.ndim - 1) + # if the input isn't MetaTensor, create MetaTensor with the default info. + data_array = convert_to_tensor(data_array, track_meta=get_track_meta()) + + original_spatial_shape = data_array.shape[1:] + sr = len(original_spatial_shape) if sr <= 0: raise ValueError("data_array must have at least one spatial dimension.") - if affine is None: + affine_: np.ndarray + affine_np: np.ndarray + if isinstance(data_array, MetaTensor): + affine_np, *_ = convert_data_type(data_array.affine, np.ndarray) + affine_ = to_affine_nd(sr, affine_np) + else: + warnings.warn("`data_array` is not of type MetaTensor, assuming affine to be identity.") # default to identity - affine_np = affine = np.eye(sr + 1, dtype=np.float64) affine_ = np.eye(sr + 1, dtype=np.float64) - else: - affine_np, *_ = convert_data_type(affine, np.ndarray) - affine_ = to_affine_nd(sr, affine_np) out_d = self.pixdim[:sr] if out_d.size < sr: @@ -432,31 +537,35 @@ def __call__( new_affine = zoom_affine(affine_, out_d, diagonal=self.diagonal) output_shape, offset = compute_shape_offset(data_array.shape[1:], affine_, new_affine) new_affine[:sr, -1] = offset[:sr] - output_data, new_affine = self.sp_resample( + # convert to MetaTensor if necessary + data_array = convert_to_tensor(data_array, track_meta=get_track_meta()) + data_array.affine = torch.as_tensor(affine_) # type: ignore + + # we don't want to track the nested transform otherwise two will be appended + data_array = self.sp_resample( data_array, - src_affine=affine, - dst_affine=new_affine, + dst_affine=torch.as_tensor(new_affine), spatial_size=list(output_shape) if output_spatial_shape is None else output_spatial_shape, mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype, ) - new_affine = to_affine_nd(affine_np, new_affine) - new_affine, *_ = convert_to_dst_type(src=new_affine, dst=affine, dtype=torch.float32) - if self.image_only: - return output_data - return output_data, affine, new_affine + return data_array + def inverse(self, data: torch.Tensor) -> torch.Tensor: + return self.sp_resample.inverse(data) -class Orientation(Transform): + +class Orientation(InvertibleTransform): """ Change the input image's orientation into the specified based on `axcodes`. """ backend = [TransformBackends.NUMPY, TransformBackends.TORCH] + @deprecated_arg(name="image_only", since="0.9") def __init__( self, axcodes: Optional[str] = None, @@ -475,7 +584,6 @@ def __init__( labels: optional, None or sequence of (2,) sequences (2,) sequences are labels for (beginning, end) of output axis. Defaults to ``(('L', 'R'), ('P', 'A'), ('I', 'S'))``. - image_only: if True return only the image volume, otherwise return (image, affine, new_affine). Raises: ValueError: When ``axcodes=None`` and ``as_closest_canonical=True``. Incompatible values. @@ -490,39 +598,41 @@ def __init__( self.axcodes = axcodes self.as_closest_canonical = as_closest_canonical self.labels = labels - self.image_only = image_only - def __call__( - self, data_array: NdarrayOrTensor, affine: Optional[NdarrayOrTensor] = None - ) -> Union[NdarrayOrTensor, Tuple[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor]]: + def __call__(self, data_array: torch.Tensor) -> torch.Tensor: """ - original orientation of `data_array` is defined by `affine`. + If input type is `MetaTensor`, original affine is extracted with `data_array.affine`. + If input type is `torch.Tensor`, original affine is assumed to be identity. Args: data_array: in shape (num_channels, H[, W, ...]). - affine (matrix): (N+1)x(N+1) original affine matrix for spatially ND `data_array`. Defaults to identity. Raises: ValueError: When ``data_array`` has no spatial dimensions. ValueError: When ``axcodes`` spatiality differs from ``data_array``. Returns: - data_array [reoriented in `self.axcodes`] if `self.image_only`, else - (data_array [reoriented in `self.axcodes`], original axcodes, current axcodes). + data_array [reoriented in `self.axcodes`]. Output type will be `MetaTensor` + unless `get_track_meta() == False`, in which case it will be + `torch.Tensor`. """ + data_array = convert_to_tensor(data_array, track_meta=get_track_meta()) + spatial_shape = data_array.shape[1:] sr = len(spatial_shape) if sr <= 0: raise ValueError("data_array must have at least one spatial dimension.") affine_: np.ndarray - if affine is None: + affine_np: np.ndarray + if isinstance(data_array, MetaTensor): + affine_np, *_ = convert_data_type(data_array.affine, np.ndarray) + affine_ = to_affine_nd(sr, affine_np) + else: + warnings.warn("`data_array` is not of type `MetaTensor, assuming affine to be identity.") # default to identity - affine_np = affine = np.eye(sr + 1, dtype=np.float64) + affine_np = np.eye(sr + 1, dtype=np.float64) affine_ = np.eye(sr + 1, dtype=np.float64) - else: - affine_np, *_ = convert_data_type(affine, np.ndarray) - affine_ = to_affine_nd(sr, affine_np) src = nib.io_orientation(affine_) if self.as_closest_canonical: @@ -543,35 +653,47 @@ def __call__( ) spatial_ornt = nib.orientations.ornt_transform(src, dst) new_affine = affine_ @ nib.orientations.inv_ornt_aff(spatial_ornt, spatial_shape) - _is_tensor = isinstance(data_array, torch.Tensor) + spatial_ornt[:, 0] += 1 # skip channel dim spatial_ornt = np.concatenate([np.array([[0, 1]]), spatial_ornt]) axes = [ax for ax, flip in enumerate(spatial_ornt[:, 1]) if flip == -1] if axes: - data_array = ( - torch.flip(data_array, dims=axes) if _is_tensor else np.flip(data_array, axis=axes) # type: ignore - ) + data_array = torch.flip(data_array, dims=axes) full_transpose = np.arange(len(data_array.shape)) full_transpose[: len(spatial_ornt)] = np.argsort(spatial_ornt[:, 0]) if not np.all(full_transpose == np.arange(len(data_array.shape))): - if _is_tensor: - data_array = data_array.permute(full_transpose.tolist()) # type: ignore - else: - data_array = data_array.transpose(full_transpose) # type: ignore - out, *_ = convert_to_dst_type(src=data_array, dst=data_array) + data_array = data_array.permute(full_transpose.tolist()) + new_affine = to_affine_nd(affine_np, new_affine) - new_affine, *_ = convert_to_dst_type(src=new_affine, dst=affine, dtype=torch.float32) + new_affine, *_ = convert_data_type(new_affine, torch.Tensor, dtype=torch.float32, device=data_array.device) + + data_array = convert_to_tensor(data_array, track_meta=get_track_meta()) + if get_track_meta(): + self.update_meta(data_array, new_affine) + self.push_transform(data_array, extra_info={"original_affine": affine_np}) + return data_array + + def update_meta(self, img, new_affine): + img.affine = new_affine + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + # Create inverse transform + orig_affine = transform[TraceKeys.EXTRA_INFO]["original_affine"] + orig_axcodes = nib.orientations.aff2axcodes(orig_affine) + inverse_transform = Orientation(axcodes=orig_axcodes, as_closest_canonical=False, labels=self.labels) + # Apply inverse + with inverse_transform.trace_transform(False): + data = inverse_transform(data) - if self.image_only: - return out - return out, affine, new_affine + return data -class Flip(Transform): +class Flip(InvertibleTransform): """ Reverses the order of elements along the given spatial axis. Preserves shape. - Uses ``np.flip`` in practice. See numpy.flip for additional details: - https://docs.scipy.org/doc/numpy/reference/generated/numpy.flip.html. + See `torch.flip` documentation for additional details: + https://pytorch.org/docs/stable/generated/torch.flip.html Args: spatial_axis: spatial axes along which to flip over. Default is None. @@ -582,22 +704,44 @@ class Flip(Transform): """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + backend = [TransformBackends.TORCH] def __init__(self, spatial_axis: Optional[Union[Sequence[int], int]] = None) -> None: self.spatial_axis = spatial_axis - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def update_meta(self, img, shape, axes): + # shape and axes include the channel dim + affine = img.affine + mat = convert_to_dst_type(torch.eye(len(affine)), affine)[0] + for axis in axes: + sp = axis - 1 + mat[sp, sp], mat[sp, -1] = mat[sp, sp] * -1, shape[axis] - 1 + img.affine = affine @ mat + + def forward_image(self, img, axes) -> torch.Tensor: + return torch.flip(img, axes) + + def __call__(self, img: torch.Tensor) -> torch.Tensor: """ Args: - img: channel first array, must have shape: (num_channels, H[, W, ..., ]), - """ - if isinstance(img, np.ndarray): - return np.ascontiguousarray(np.flip(img, map_spatial_axes(img.ndim, self.spatial_axis))) - return torch.flip(img, map_spatial_axes(img.ndim, self.spatial_axis)) + img: channel first array, must have shape: (num_channels, H[, W, ..., ]) + """ + img = convert_to_tensor(img, track_meta=get_track_meta()) + axes = map_spatial_axes(img.ndim, self.spatial_axis) + out = self.forward_image(img, axes) + if get_track_meta(): + self.update_meta(out, out.shape, axes) + self.push_transform(out) + return out + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + self.pop_transform(data) + flipper = Flip(spatial_axis=self.spatial_axis) + with flipper.trace_transform(False): + return flipper(data) -class Resize(Transform): +class Resize(InvertibleTransform): """ Resize the input image to given spatial size (with scaling, not cropping/padding). Implemented using :py:class:`torch.nn.functional.interpolate`. @@ -612,12 +756,21 @@ class Resize(Transform): which must be an int number in this case, keeping the aspect ratio of the initial image, refer to: https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/ #albumentations.augmentations.geometric.resize.LongestMaxSize. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html align_corners: This only has an effect when mode is 'linear', 'bilinear', 'bicubic' or 'trilinear'. Default: None. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + anti_aliasing: bool + Whether to apply a Gaussian filter to smooth the image prior + to downsampling. It is crucial to filter when downsampling + the image to avoid aliasing artifacts. See also ``skimage.transform.resize`` + anti_aliasing_sigma: {float, tuple of floats}, optional + Standard deviation for Gaussian filtering used when anti-aliasing. + By default, this value is chosen as (s - 1) / 2 where s is the + downsampling factor, where s > 1. For the up-size case, s < 1, no + anti-aliasing is performed prior to rescaling. """ backend = [TransformBackends.TORCH] @@ -626,64 +779,135 @@ def __init__( self, spatial_size: Union[Sequence[int], int], size_mode: str = "all", - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, + mode: str = InterpolateMode.AREA, align_corners: Optional[bool] = None, + anti_aliasing: bool = False, + anti_aliasing_sigma: Union[Sequence[float], float, None] = None, ) -> None: self.size_mode = look_up_option(size_mode, ["all", "longest"]) self.spatial_size = spatial_size self.mode: InterpolateMode = look_up_option(mode, InterpolateMode) self.align_corners = align_corners + self.anti_aliasing = anti_aliasing + self.anti_aliasing_sigma = anti_aliasing_sigma def __call__( self, - img: NdarrayOrTensor, - mode: Optional[Union[InterpolateMode, str]] = None, + img: torch.Tensor, + mode: Optional[str] = None, align_corners: Optional[bool] = None, - ) -> NdarrayOrTensor: + anti_aliasing: Optional[bool] = None, + anti_aliasing_sigma: Union[Sequence[float], float, None] = None, + ) -> torch.Tensor: """ Args: img: channel first array, must have shape: (num_channels, H[, W, ..., ]). - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, + ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``self.mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html align_corners: This only has an effect when mode is 'linear', 'bilinear', 'bicubic' or 'trilinear'. Defaults to ``self.align_corners``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html + anti_aliasing: bool, optional + Whether to apply a Gaussian filter to smooth the image prior + to downsampling. It is crucial to filter when downsampling + the image to avoid aliasing artifacts. See also ``skimage.transform.resize`` + anti_aliasing_sigma: {float, tuple of floats}, optional + Standard deviation for Gaussian filtering used when anti-aliasing. + By default, this value is chosen as (s - 1) / 2 where s is the + downsampling factor, where s > 1. For the up-size case, s < 1, no + anti-aliasing is performed prior to rescaling. Raises: ValueError: When ``self.spatial_size`` length is less than ``img`` spatial dimensions. """ - img_, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float) + 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 if self.size_mode == "all": - input_ndim = img_.ndim - 1 # spatial ndim 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) + input_shape = ensure_tuple_size(img.shape, output_ndim + 1, 1) + img = img.reshape(input_shape) elif output_ndim < input_ndim: raise ValueError( "len(spatial_size) must be greater or equal to img spatial dimensions, " f"got spatial_size={output_ndim} img={input_ndim}." ) - spatial_size_ = fall_back_tuple(self.spatial_size, img_.shape[1:]) + spatial_size_ = fall_back_tuple(self.spatial_size, img.shape[1:]) else: # for the "longest" mode - img_size = img_.shape[1:] + img_size = img.shape[1:] if not isinstance(self.spatial_size, int): raise ValueError("spatial_size must be an int number if size_mode is 'longest'.") scale = self.spatial_size / max(img_size) spatial_size_ = tuple(int(round(s * scale)) for s in img_size) + + if tuple(img.shape[1:]) == spatial_size_: # spatial shape is already the desired + return convert_to_tensor(img, track_meta=get_track_meta()) # type: ignore + + original_sp_size = img.shape[1:] + img_ = convert_to_tensor(img, dtype=torch.float, track_meta=False) + + if anti_aliasing and any(x < y for x, y in zip(spatial_size_, img_.shape[1:])): + factors = torch.div(torch.Tensor(list(img_.shape[1:])), torch.Tensor(spatial_size_)) + if anti_aliasing_sigma is None: + # if sigma is not given, use the default sigma in skimage.transform.resize + anti_aliasing_sigma = torch.maximum(torch.zeros(factors.shape), (factors - 1) / 2).tolist() + else: + # if sigma is given, use the given value for downsampling axis + anti_aliasing_sigma = list(ensure_tuple_rep(anti_aliasing_sigma, len(spatial_size_))) + for axis in range(len(spatial_size_)): + anti_aliasing_sigma[axis] = anti_aliasing_sigma[axis] * int(factors[axis] > 1) + anti_aliasing_filter = GaussianSmooth(sigma=anti_aliasing_sigma) + img_ = convert_to_tensor(anti_aliasing_filter(img_), track_meta=False) + + img = convert_to_tensor(img, track_meta=get_track_meta()) + _mode = look_up_option(self.mode if mode is None else mode, InterpolateMode) + _align_corners = self.align_corners if align_corners is None else align_corners + resized = torch.nn.functional.interpolate( - input=img_.unsqueeze(0), - size=spatial_size_, - mode=look_up_option(self.mode if mode is None else mode, InterpolateMode).value, - align_corners=self.align_corners if align_corners is None else align_corners, + input=img_.unsqueeze(0), size=spatial_size_, mode=_mode, align_corners=_align_corners ) out, *_ = convert_to_dst_type(resized.squeeze(0), img) + if get_track_meta(): + self.update_meta(out, original_sp_size, spatial_size_) + self.push_transform( + out, + orig_size=original_sp_size, + extra_info={ + "mode": _mode, + "align_corners": _align_corners if _align_corners is not None else TraceKeys.NONE, + "new_dim": len(original_sp_size) - input_ndim, # additional dims appended + }, + ) return out + def update_meta(self, img, spatial_size, new_spatial_size): + affine = convert_to_tensor(img.affine, track_meta=False) + img.affine = scale_affine(affine, spatial_size, new_spatial_size) + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + return self.inverse_transform(data, transform) + + def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor: + orig_size = transform[TraceKeys.ORIG_SIZE] + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] + xform = Resize( + spatial_size=orig_size, mode=mode, align_corners=None if align_corners == TraceKeys.NONE else align_corners + ) + with xform.trace_transform(False): + data = xform(data) + for _ in range(transform[TraceKeys.EXTRA_INFO]["new_dim"]): + data = data.squeeze(-1) # remove the additional dims + return data -class Rotate(Transform, ThreadUnsafe): + +class Rotate(InvertibleTransform): """ Rotates an input image by given angle using :py:class:`monai.networks.layers.AffineTransform`. @@ -711,27 +935,26 @@ def __init__( self, angle: Union[Sequence[float], float], keep_size: bool = True, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: Union[DtypeLike, torch.dtype] = np.float32, ) -> None: self.angle = angle self.keep_size = keep_size - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) self.align_corners = align_corners self.dtype = dtype - self._rotation_matrix: Optional[NdarrayOrTensor] = None def __call__( self, - img: NdarrayOrTensor, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + img: torch.Tensor, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = None, dtype: Union[DtypeLike, torch.dtype] = None, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: channel first array, must have shape: [chns, H, W] or [chns, H, W, D]. @@ -753,14 +976,13 @@ def __call__( ValueError: When ``img`` spatially is not one of [2D, 3D]. """ - _dtype = dtype or self.dtype or img.dtype + img = convert_to_tensor(img, track_meta=get_track_meta()) + _dtype = get_equivalent_dtype(dtype or self.dtype or img.dtype, torch.Tensor) - img_t, *_ = convert_data_type(img, torch.Tensor, dtype=_dtype) - - im_shape = np.asarray(img_t.shape[1:]) # spatial dimensions + im_shape = np.asarray(img.shape[1:]) # spatial dimensions input_ndim = len(im_shape) if input_ndim not in (2, 3): - raise ValueError(f"Unsupported img dimension: {input_ndim}, available options are [2, 3].") + raise ValueError(f"Unsupported image dimension: {input_ndim}, available options are [2, 3].") _angle = ensure_tuple_rep(self.angle, 1 if input_ndim == 2 else 3) transform = create_rotate(input_ndim, _angle) shift = create_translate(input_ndim, ((im_shape - 1) / 2).tolist()) @@ -775,30 +997,70 @@ def __call__( shift_1 = create_translate(input_ndim, (-(output_shape - 1) / 2).tolist()) transform = shift @ transform @ shift_1 + img_t = img.to(_dtype) transform_t, *_ = convert_to_dst_type(transform, img_t) - + _mode = look_up_option(mode or self.mode, GridSampleMode) + _padding_mode = look_up_option(padding_mode or self.padding_mode, GridSamplePadMode) + _align_corners = self.align_corners if align_corners is None else align_corners xform = AffineTransform( normalized=False, - mode=look_up_option(mode or self.mode, GridSampleMode), - padding_mode=look_up_option(padding_mode or self.padding_mode, GridSamplePadMode), - align_corners=self.align_corners if align_corners is None else align_corners, + mode=_mode, + padding_mode=_padding_mode, + align_corners=_align_corners, reverse_indexing=True, ) output: torch.Tensor = xform(img_t.unsqueeze(0), transform_t, spatial_size=output_shape).float().squeeze(0) - self._rotation_matrix = transform - out: NdarrayOrTensor out, *_ = convert_to_dst_type(output, dst=img, dtype=output.dtype) + if get_track_meta(): + self.update_meta(out, transform_t) + self.push_transform( + out, + orig_size=img_t.shape[1:], + extra_info={ + "rot_mat": transform, + "mode": _mode, + "padding_mode": _padding_mode, + "align_corners": _align_corners if _align_corners is not None else TraceKeys.NONE, + "dtype": str(_dtype)[6:], # dtype as string; remove "torch": torch.float32 -> float32 + }, + ) return out - def get_rotation_matrix(self) -> Optional[NdarrayOrTensor]: - """ - Get the most recently applied rotation matrix - This is not thread-safe. - """ - return self._rotation_matrix + def update_meta(self, img, rotate_mat): + affine = convert_to_tensor(img.affine, track_meta=False) + mat = to_affine_nd(len(affine) - 1, rotate_mat) + img.affine = affine @ convert_to_dst_type(mat, affine)[0] + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + return self.inverse_transform(data, transform) + + def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor: + fwd_rot_mat = transform[TraceKeys.EXTRA_INFO]["rot_mat"] + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] + align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] + dtype = transform[TraceKeys.EXTRA_INFO]["dtype"] + inv_rot_mat = linalg_inv(fwd_rot_mat) + + xform = AffineTransform( + normalized=False, + mode=mode, + padding_mode=padding_mode, + align_corners=False if align_corners == TraceKeys.NONE else align_corners, + reverse_indexing=True, + ) + img_t: torch.Tensor = convert_data_type(data, MetaTensor, dtype=dtype)[0] + transform_t, *_ = convert_to_dst_type(inv_rot_mat, img_t) + sp_size = transform[TraceKeys.ORIG_SIZE] + out: torch.Tensor = xform(img_t.unsqueeze(0), transform_t, spatial_size=sp_size).float().squeeze(0) + out = convert_to_dst_type(out, dst=data, dtype=out.dtype)[0] + if isinstance(data, MetaTensor): + self.update_meta(out, transform_t) + return out -class Zoom(Transform): +class Zoom(InvertibleTransform): """ Zooms an ND image using :py:class:`torch.nn.functional.interpolate`. For details, please see https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html. @@ -810,7 +1072,7 @@ class Zoom(Transform): zoom: The zoom factor along the spatial axes. If a float, zoom is the same for each spatial axis. If a sequence, zoom should contain one value for each spatial axis. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, @@ -834,8 +1096,8 @@ class Zoom(Transform): def __init__( self, zoom: Union[Sequence[float], float], - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, - padding_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.EDGE, + mode: str = InterpolateMode.AREA, + padding_mode: str = NumpyPadMode.EDGE, align_corners: Optional[bool] = None, keep_size: bool = True, **kwargs, @@ -849,15 +1111,16 @@ def __init__( def __call__( self, - img: NdarrayOrTensor, - mode: Optional[Union[InterpolateMode, str]] = None, - padding_mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None, + img: torch.Tensor, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = None, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: channel first array, must have shape: (num_channels, H[, W, ..., ]). - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, + ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``self.mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, @@ -872,47 +1135,83 @@ def __call__( See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html """ - img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float32) + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_t = img.to(torch.float32) _zoom = ensure_tuple_rep(self.zoom, img.ndim - 1) # match the spatial image dim - zoomed: NdarrayOrTensor = torch.nn.functional.interpolate( # type: ignore + _mode = look_up_option(self.mode if mode is None else mode, InterpolateMode).value + _align_corners = self.align_corners if align_corners is None else align_corners + _padding_mode = padding_mode or self.padding_mode + + zoomed: NdarrayOrTensor = torch.nn.functional.interpolate( recompute_scale_factor=True, input=img_t.unsqueeze(0), scale_factor=list(_zoom), - mode=look_up_option(self.mode if mode is None else mode, InterpolateMode).value, - align_corners=self.align_corners if align_corners is None else align_corners, + mode=_mode, + align_corners=_align_corners, ) zoomed = zoomed.squeeze(0) - - if self.keep_size and not np.allclose(img_t.shape, zoomed.shape): - - pad_vec = [(0, 0)] * len(img_t.shape) - slice_vec = [slice(None)] * len(img_t.shape) - for idx, (od, zd) in enumerate(zip(img_t.shape, zoomed.shape)): - diff = od - zd - half = abs(diff) // 2 - if diff > 0: # need padding - pad_vec[idx] = (half, diff - half) - elif diff < 0: # need slicing - slice_vec[idx] = slice(half, half + od) - - padder = Pad(pad_vec, padding_mode or self.padding_mode) - zoomed = padder(zoomed) - zoomed = zoomed[tuple(slice_vec)] + orig_size, z_size = img_t.shape, zoomed.shape out, *_ = convert_to_dst_type(zoomed, dst=img) + if get_track_meta(): + self.update_meta(out, orig_size[1:], z_size[1:]) + do_pad_crop = self.keep_size and not np.allclose(orig_size, z_size) + if do_pad_crop: + _pad_crop = ResizeWithPadOrCrop(spatial_size=img_t.shape[1:], mode=_padding_mode) + out = _pad_crop(out) + if get_track_meta(): + padcrop_xform = self.pop_transform(out, check=False) if do_pad_crop else {} + self.push_transform( + out, + orig_size=orig_size[1:], + extra_info={ + "mode": _mode, + "align_corners": _align_corners if _align_corners is not None else TraceKeys.NONE, + "do_padcrop": do_pad_crop, + "padcrop": padcrop_xform, + }, + ) + return out + + def update_meta(self, img, spatial_size, new_spatial_size): + affine = convert_to_tensor(img.affine, track_meta=False) + img.affine = scale_affine(affine, spatial_size, new_spatial_size) + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + return self.inverse_transform(data, transform) + + def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor: + if transform[TraceKeys.EXTRA_INFO]["do_padcrop"]: + orig_size = transform[TraceKeys.ORIG_SIZE] + pad_or_crop = ResizeWithPadOrCrop(spatial_size=orig_size, mode="edge") + padcrop_xform = transform[TraceKeys.EXTRA_INFO]["padcrop"] + padcrop_xform[TraceKeys.EXTRA_INFO]["pad_info"][TraceKeys.ID] = TraceKeys.NONE + padcrop_xform[TraceKeys.EXTRA_INFO]["crop_info"][TraceKeys.ID] = TraceKeys.NONE + # this uses inverse because spatial_size // 2 in the forward pass of center crop may cause issues + data = pad_or_crop.inverse_transform(data, padcrop_xform) # type: ignore + # Create inverse transform + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] + inverse_transform = Resize(spatial_size=transform[TraceKeys.ORIG_SIZE]) + # Apply inverse + with inverse_transform.trace_transform(False): + out = inverse_transform( + data, mode=mode, align_corners=None if align_corners == TraceKeys.NONE else align_corners + ) return out -class Rotate90(Transform): +class Rotate90(InvertibleTransform): """ Rotate an array by 90 degrees in the plane specified by `axes`. - See np.rot90 for additional details: - https://numpy.org/doc/stable/reference/generated/numpy.rot90.html. + See `torch.rot90` for additional details: + https://pytorch.org/docs/stable/generated/torch.rot90.html#torch-rot90. """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + backend = [TransformBackends.TORCH] def __init__(self, k: int = 1, spatial_axes: Tuple[int, int] = (0, 1)) -> None: """ @@ -928,18 +1227,52 @@ def __init__(self, k: int = 1, spatial_axes: Tuple[int, int] = (0, 1)) -> None: raise ValueError("spatial_axes must be 2 int numbers to indicate the axes to rotate 90 degrees.") self.spatial_axes = spatial_axes_ - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor) -> torch.Tensor: """ Args: img: channel first array, must have shape: (num_channels, H[, W, ..., ]), """ - rot90: Callable = torch.rot90 if isinstance(img, torch.Tensor) else np.rot90 # type: ignore - out: NdarrayOrTensor = rot90(img, self.k, map_spatial_axes(img.ndim, self.spatial_axes)) - out, *_ = convert_data_type(out, dtype=img.dtype) + img = convert_to_tensor(img, track_meta=get_track_meta()) + axes = map_spatial_axes(img.ndim, self.spatial_axes) + ori_shape = img.shape[1:] + out: NdarrayOrTensor = torch.rot90(img, self.k, axes) + out = convert_to_dst_type(out, img)[0] + if get_track_meta(): + self.update_meta(out, ori_shape, out.shape[1:], axes, self.k) + self.push_transform(out, extra_info={"axes": [d - 1 for d in axes], "k": self.k}) # compensate spatial dim return out - -class RandRotate90(RandomizableTransform): + def update_meta(self, img, spatial_size, new_spatial_size, axes, k): + affine = convert_data_type(img.affine, torch.Tensor)[0] + r, sp_r = len(affine) - 1, len(spatial_size) + mat = to_affine_nd(r, create_translate(sp_r, [-float(d - 1) / 2 for d in new_spatial_size])) + s = -1.0 if int(axes[0]) - int(axes[1]) in (-1, 2) else 1.0 + if sp_r == 2: + rot90 = to_affine_nd(r, create_rotate(sp_r, [s * np.pi / 2])) + else: + idx = {1, 2, 3} - set(axes) + angle = [0, 0, 0] + angle[idx.pop() - 1] = s * np.pi / 2 + rot90 = to_affine_nd(r, create_rotate(sp_r, angle)) + for _ in range(k): + mat = rot90 @ mat + mat = to_affine_nd(r, create_translate(sp_r, [float(d - 1) / 2 for d in spatial_size])) @ mat + img.affine = affine @ convert_to_dst_type(mat, affine)[0] + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + return self.inverse_transform(data, transform) + + def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor: + axes = transform[TraceKeys.EXTRA_INFO]["axes"] + k = transform[TraceKeys.EXTRA_INFO]["k"] + inv_k = 4 - k % 4 + xform = Rotate90(k=inv_k, spatial_axes=axes) + with xform.trace_transform(False): + return xform(data) + + +class RandRotate90(RandomizableTransform, InvertibleTransform): """ With probability `prob`, input arrays are rotated by 90 degrees in the plane specified by `spatial_axes`. @@ -968,7 +1301,7 @@ def randomize(self, data: Optional[Any] = None) -> None: return None self._rand_k = self.R.randint(self.max_k) + 1 - def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor: """ Args: img: channel first array, must have shape: (num_channels, H[, W, ..., ]), @@ -977,13 +1310,25 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen if randomize: self.randomize() - if not self._do_transform: - return img + if self._do_transform: + out = Rotate90(self._rand_k, self.spatial_axes)(img) + else: + out = convert_to_tensor(img, track_meta=get_track_meta()) + + if get_track_meta(): + maybe_rot90_info = self.pop_transform(out, check=False) if self._do_transform else {} + self.push_transform(out, extra_info=maybe_rot90_info) + return out - return Rotate90(self._rand_k, self.spatial_axes)(img) + def inverse(self, data: torch.Tensor) -> torch.Tensor: + xform_info = self.pop_transform(data) + if not xform_info[TraceKeys.DO_TRANSFORM]: + return data + rotate_xform = xform_info[TraceKeys.EXTRA_INFO] + return Rotate90().inverse_transform(data, rotate_xform) -class RandRotate(RandomizableTransform): +class RandRotate(RandomizableTransform, InvertibleTransform): """ Randomly rotate the input arrays. @@ -991,9 +1336,9 @@ class RandRotate(RandomizableTransform): range_x: Range of rotation angle in radians in the plane defined by the first and second axes. If single number, angle is uniformly sampled from (-range_x, range_x). range_y: Range of rotation angle in radians in the plane defined by the first and third axes. - If single number, angle is uniformly sampled from (-range_y, range_y). + If single number, angle is uniformly sampled from (-range_y, range_y). only work for 3D data. range_z: Range of rotation angle in radians in the plane defined by the second and third axes. - If single number, angle is uniformly sampled from (-range_z, range_z). + If single number, angle is uniformly sampled from (-range_z, range_z). only work for 3D data. prob: Probability of rotation. keep_size: If it is False, the output shape is adapted so that the input array is contained completely in the output. @@ -1020,8 +1365,8 @@ def __init__( range_z: Union[Tuple[float, float], float] = 0.0, prob: float = 0.1, keep_size: bool = True, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, align_corners: bool = False, dtype: Union[DtypeLike, torch.dtype] = np.float32, ) -> None: @@ -1037,8 +1382,8 @@ def __init__( self.range_z = tuple(sorted([-self.range_z[0], self.range_z[0]])) self.keep_size = keep_size - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) self.align_corners = align_corners self.dtype = dtype @@ -1054,11 +1399,12 @@ def randomize(self, data: Optional[Any] = None) -> None: self.y = self.R.uniform(low=self.range_y[0], high=self.range_y[1]) self.z = self.R.uniform(low=self.range_z[0], high=self.range_z[1]) + @deprecated_arg(name="get_matrix", since="0.9", msg_suffix="please use `img.meta` instead.") def __call__( self, - img: NdarrayOrTensor, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + img: torch.Tensor, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = None, dtype: Union[DtypeLike, torch.dtype] = None, randomize: bool = True, @@ -1079,27 +1425,35 @@ def __call__( If None, use the data type of input data. To be compatible with other modules, the output data type is always ``np.float32``. randomize: whether to execute `randomize()` function first, default to True. - get_matrix: whether to return the rotated image and rotate matrix together, default to False. """ if randomize: self.randomize() - if not self._do_transform: - return img - - rotator = Rotate( - angle=self.x if img.ndim == 3 else (self.x, self.y, self.z), - keep_size=self.keep_size, - mode=look_up_option(mode or self.mode, GridSampleMode), - padding_mode=look_up_option(padding_mode or self.padding_mode, GridSamplePadMode), - align_corners=self.align_corners if align_corners is None else align_corners, - dtype=dtype or self.dtype or img.dtype, - ) - img = rotator(img) - return (img, rotator.get_rotation_matrix()) if get_matrix else img + if self._do_transform: + rotator = Rotate( + angle=self.x if img.ndim == 3 else (self.x, self.y, self.z), + keep_size=self.keep_size, + mode=look_up_option(mode or self.mode, GridSampleMode), + padding_mode=look_up_option(padding_mode or self.padding_mode, GridSamplePadMode), + align_corners=self.align_corners if align_corners is None else align_corners, + dtype=dtype or self.dtype or img.dtype, + ) + out = rotator(img) + else: + out = convert_to_tensor(img, track_meta=get_track_meta()) + if get_track_meta(): + rot_info = self.pop_transform(out, check=False) if self._do_transform else {} + self.push_transform(out, extra_info=rot_info) + return out + def inverse(self, data: torch.Tensor) -> torch.Tensor: + xform_info = self.pop_transform(data) + if not xform_info[TraceKeys.DO_TRANSFORM]: + return data + return Rotate(0).inverse_transform(data, xform_info[TraceKeys.EXTRA_INFO]) -class RandFlip(RandomizableTransform): + +class RandFlip(RandomizableTransform, InvertibleTransform): """ Randomly flips the image along axes. Preserves shape. See numpy.flip for additional details. @@ -1116,7 +1470,7 @@ def __init__(self, prob: float = 0.1, spatial_axis: Optional[Union[Sequence[int] RandomizableTransform.__init__(self, prob) self.flipper = Flip(spatial_axis=spatial_axis) - def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor: """ Args: img: channel first array, must have shape: (num_channels, H[, W, ..., ]), @@ -1124,14 +1478,22 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen """ if randomize: self.randomize(None) + out = self.flipper(img) if self._do_transform else img + out = convert_to_tensor(out, track_meta=get_track_meta()) + if get_track_meta(): + xform_info = self.pop_transform(out, check=False) if self._do_transform else {} + self.push_transform(out, extra_info=xform_info) + return out - if not self._do_transform: - return img - - return self.flipper(img) + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + if not transform[TraceKeys.DO_TRANSFORM]: + return data + data.applied_operations.append(transform[TraceKeys.EXTRA_INFO]) # type: ignore + return self.flipper.inverse(data) -class RandAxisFlip(RandomizableTransform): +class RandAxisFlip(RandomizableTransform, InvertibleTransform): """ Randomly select a spatial axis and flip along it. See numpy.flip for additional details. @@ -1147,6 +1509,7 @@ class RandAxisFlip(RandomizableTransform): def __init__(self, prob: float = 0.1) -> None: RandomizableTransform.__init__(self, prob) self._axis: Optional[int] = None + self.flipper = Flip(spatial_axis=self._axis) def randomize(self, data: NdarrayOrTensor) -> None: super().randomize(None) @@ -1154,22 +1517,36 @@ def randomize(self, data: NdarrayOrTensor) -> None: return None self._axis = self.R.randint(data.ndim - 1) - def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor: """ Args: - img: channel first array, must have shape: (num_channels, H[, W, ..., ]), + img: channel first array, must have shape: (num_channels, H[, W, ..., ]) randomize: whether to execute `randomize()` function first, default to True. """ if randomize: self.randomize(data=img) - if not self._do_transform: - return img + if self._do_transform: + self.flipper.spatial_axis = self._axis + out = self.flipper(img) + else: + out = convert_to_tensor(img, track_meta=get_track_meta()) + if get_track_meta(): + xform = self.pop_transform(out, check=False) if self._do_transform else {} + xform["axes"] = self._axis + self.push_transform(out, extra_info=xform) + return out - return Flip(spatial_axis=self._axis)(img) + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + if not transform[TraceKeys.DO_TRANSFORM]: + return data + flipper = Flip(spatial_axis=transform[TraceKeys.EXTRA_INFO]["axes"]) + with flipper.trace_transform(False): + return flipper(data) -class RandZoom(RandomizableTransform): +class RandZoom(RandomizableTransform, InvertibleTransform): """ Randomly zooms input arrays with given probability within given zoom range. @@ -1185,7 +1562,7 @@ class RandZoom(RandomizableTransform): to keep the original spatial shape ratio. If a sequence, max_zoom should contain one value for each spatial axis. If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, @@ -1211,8 +1588,8 @@ def __init__( prob: float = 0.1, min_zoom: Union[Sequence[float], float] = 0.9, max_zoom: Union[Sequence[float], float] = 1.1, - mode: Union[InterpolateMode, str] = InterpolateMode.AREA, - padding_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.EDGE, + mode: str = InterpolateMode.AREA, + padding_mode: str = NumpyPadMode.EDGE, align_corners: Optional[bool] = None, keep_size: bool = True, **kwargs, @@ -1244,17 +1621,17 @@ def randomize(self, img: NdarrayOrTensor) -> None: def __call__( self, - img: NdarrayOrTensor, - mode: Optional[Union[InterpolateMode, str]] = None, - padding_mode: Optional[Union[NumpyPadMode, PytorchPadMode, str]] = None, + img: torch.Tensor, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, align_corners: Optional[bool] = None, randomize: bool = True, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: channel first array, must have shape 2D: (nchannels, H, W), or 3D: (nchannels, H, W, D). - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} - The interpolation mode. Defaults to ``self.mode``. + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, + ``"area"``}, the interpolation mode. Defaults to ``self.mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} @@ -1274,16 +1651,26 @@ def __call__( self.randomize(img=img) if not self._do_transform: - return img - - return Zoom( - self._zoom, - keep_size=self.keep_size, - mode=look_up_option(mode or self.mode, InterpolateMode), - padding_mode=padding_mode or self.padding_mode, - align_corners=self.align_corners if align_corners is None else align_corners, - **self.kwargs, - )(img) + out = convert_to_tensor(img, track_meta=get_track_meta()) + else: + out = Zoom( + self._zoom, + keep_size=self.keep_size, + mode=look_up_option(mode or self.mode, InterpolateMode), + padding_mode=padding_mode or self.padding_mode, + align_corners=self.align_corners if align_corners is None else align_corners, + **self.kwargs, + )(img) + if get_track_meta(): + z_info = self.pop_transform(out, check=False) if self._do_transform else {} + self.push_transform(out, extra_info=z_info) + return out # type: ignore + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + xform_info = self.pop_transform(data) + if not xform_info[TraceKeys.DO_TRANSFORM]: + return data + return Zoom(self._zoom).inverse_transform(data, xform_info[TraceKeys.EXTRA_INFO]) class AffineGrid(Transform): @@ -1319,7 +1706,7 @@ class AffineGrid(Transform): """ - backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + backend = [TransformBackends.TORCH] @deprecated_arg(name="as_tensor_output", since="0.6") def __init__( @@ -1342,8 +1729,8 @@ def __init__( self.affine = affine def __call__( - self, spatial_size: Optional[Sequence[int]] = None, grid: Optional[NdarrayOrTensor] = None - ) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: + self, spatial_size: Optional[Sequence[int]] = None, grid: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: """ The grid can be initialized with a `spatial_size` parameter, or provided directly as `grid`. Therefore, either `spatial_size` or `grid` must be provided. @@ -1360,17 +1747,17 @@ def __call__( if grid is None: # create grid from spatial_size if spatial_size is None: raise ValueError("Incompatible values: grid=None and spatial_size=None.") - grid = create_grid(spatial_size, device=self.device, backend="torch", dtype=self.dtype) - _b = TransformBackends.TORCH if isinstance(grid, torch.Tensor) else TransformBackends.NUMPY - _device = grid.device if isinstance(grid, torch.Tensor) else self.device + grid_ = create_grid(spatial_size, device=self.device, backend="torch", dtype=self.dtype) + else: + grid_ = grid + _dtype = self.dtype or grid_.dtype + grid_: torch.Tensor = convert_to_tensor(grid_, dtype=_dtype, track_meta=get_track_meta()) # type: ignore + _b = TransformBackends.TORCH + _device = grid_.device # type: ignore affine: NdarrayOrTensor if self.affine is None: - spatial_dims = len(grid.shape) - 1 - affine = ( - torch.eye(spatial_dims + 1, device=_device) - if _b == TransformBackends.TORCH - else np.eye(spatial_dims + 1) - ) + spatial_dims = len(grid_.shape) - 1 + affine = torch.eye(spatial_dims + 1, device=_device) if self.rotate_params: affine = affine @ create_rotate(spatial_dims, self.rotate_params, device=_device, backend=_b) if self.shear_params: @@ -1382,11 +1769,10 @@ def __call__( else: affine = self.affine - grid, *_ = convert_data_type(grid, torch.Tensor, device=_device, dtype=self.dtype or grid.dtype) - affine, *_ = convert_to_dst_type(affine, grid) - - grid = (affine @ grid.reshape((grid.shape[0], -1))).reshape([-1] + list(grid.shape[1:])) - return grid, affine + affine = to_affine_nd(len(grid_) - 1, affine) + affine = convert_to_tensor(affine, device=grid_.device, dtype=grid_.dtype, track_meta=False) # type: ignore + grid_ = (affine @ grid_.reshape((grid_.shape[0], -1))).reshape([-1] + list(grid_.shape[1:])) + return grid_, affine # type: ignore class RandAffineGrid(Randomizable, Transform): @@ -1454,7 +1840,7 @@ def __init__( self.scale_params: Optional[List[float]] = None self.device = device - self.affine: Optional[NdarrayOrTensor] = None + self.affine: Optional[torch.Tensor] = torch.eye(4, dtype=torch.float64) def _get_rand_param(self, param_range, add_scalar: float = 0.0): out_param = [] @@ -1474,17 +1860,22 @@ def randomize(self, data: Optional[Any] = None) -> None: self.scale_params = self._get_rand_param(self.scale_range, 1.0) def __call__( - self, spatial_size: Optional[Sequence[int]] = None, grid: Optional[NdarrayOrTensor] = None - ) -> NdarrayOrTensor: + self, + spatial_size: Optional[Sequence[int]] = None, + grid: Optional[NdarrayOrTensor] = None, + randomize: bool = True, + ) -> torch.Tensor: """ Args: spatial_size: output grid size. grid: grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D. + randomize: boolean as to whether the grid parameters governing the grid should be randomized. Returns: a 2D (3xHxW) or 3D (4xHxWxD) grid. """ - self.randomize() + if randomize: + self.randomize() affine_grid = AffineGrid( rotate_params=self.rotate_params, shear_params=self.shear_params, @@ -1492,11 +1883,11 @@ def __call__( scale_params=self.scale_params, device=self.device, ) - _grid: NdarrayOrTensor + _grid: torch.Tensor _grid, self.affine = affine_grid(spatial_size, grid) return _grid - def get_transformation_matrix(self) -> Optional[NdarrayOrTensor]: + def get_transformation_matrix(self) -> Optional[torch.Tensor]: """Get the most recently applied transformation matrix""" return self.affine @@ -1508,6 +1899,7 @@ class RandDeformGrid(Randomizable, Transform): backend = [TransformBackends.TORCH] + @deprecated_arg(name="as_tensor_output", since="0.8") def __init__( self, spacing: Union[Sequence[float], float], @@ -1523,15 +1915,12 @@ def __init__( spacing=(2, 2) indicates deformation field defined on every other pixel in 2D. magnitude_range: the random offsets will be generated from `uniform[magnitude[0], magnitude[1])`. - as_tensor_output: whether to output tensor instead of numpy array. - defaults to True. device: device to store the output grid data. """ self.spacing = spacing self.magnitude = magnitude_range self.rand_mag = 1.0 - self.as_tensor_output = as_tensor_output self.random_offset: np.ndarray self.device = device @@ -1539,7 +1928,7 @@ def randomize(self, grid_size: Sequence[int]) -> None: self.random_offset = self.R.normal(size=([len(grid_size)] + list(grid_size))).astype(np.float32, copy=False) self.rand_mag = self.R.uniform(self.magnitude[0], self.magnitude[1]) - def __call__(self, spatial_size: Sequence[int]): + def __call__(self, spatial_size: Sequence[int]) -> torch.Tensor: """ Args: spatial_size: spatial size of the grid. @@ -1549,9 +1938,7 @@ def __call__(self, spatial_size: Sequence[int]): self.randomize(control_grid.shape[1:]) _offset, *_ = convert_to_dst_type(self.rand_mag * self.random_offset, control_grid) control_grid[: len(spatial_size)] += _offset - if not self.as_tensor_output: - control_grid, *_ = convert_data_type(control_grid, output_type=np.ndarray, dtype=np.float32) - return control_grid + return control_grid # type: ignore class Resample(Transform): @@ -1561,8 +1948,8 @@ class Resample(Transform): @deprecated_arg(name="as_tensor_output", since="0.6") def __init__( self, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, as_tensor_output: bool = True, norm_coords: bool = True, device: Optional[torch.device] = None, @@ -1595,20 +1982,20 @@ def __init__( ``as_tensor_output`` is deprecated. """ - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) self.norm_coords = norm_coords self.device = device self.dtype = dtype - def __call__( + def __call__( # type: ignore self, - img: NdarrayOrTensor, - grid: Optional[NdarrayOrTensor] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + img: torch.Tensor, + grid: torch.Tensor, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, dtype: DtypeLike = None, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W[, D]). @@ -1631,12 +2018,11 @@ def __call__( See also: :py:const:`monai.config.USE_COMPILED` """ - if grid is None: - raise ValueError("Unknown grid.") _device = img.device if isinstance(img, torch.Tensor) else self.device _dtype = dtype or self.dtype or img.dtype - img_t, *_ = convert_data_type(img, torch.Tensor, device=_device, dtype=_dtype) - grid_t = convert_to_dst_type(grid, img_t)[0] + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_t, *_ = convert_data_type(img, torch.Tensor, dtype=_dtype, device=_device) + grid_t, *_ = convert_to_dst_type(grid, img_t) if grid_t is grid: # copy if needed (convert_data_type converts to contiguous) grid_t = grid_t.clone(memory_format=torch.contiguous_format) sr = min(len(img_t.shape[1:]), 3) @@ -1647,10 +2033,8 @@ def __call__( grid_t[i] = (max(dim, 2) / 2.0 - 0.5 + grid_t[i]) / grid_t[-1:] grid_t = moveaxis(grid_t[:sr], 0, -1) # type: ignore _padding_mode = self.padding_mode if padding_mode is None else padding_mode - _padding_mode = _padding_mode.value if isinstance(_padding_mode, GridSamplePadMode) else _padding_mode bound = 1 if _padding_mode == "reflection" else _padding_mode _interp_mode = self.mode if mode is None else mode - _interp_mode = _interp_mode.value if isinstance(_interp_mode, GridSampleMode) else _interp_mode if _interp_mode == "bicubic": interp = 3 elif _interp_mode == "bilinear": @@ -1669,15 +2053,15 @@ def __call__( out = torch.nn.functional.grid_sample( img_t.unsqueeze(0), grid_t.unsqueeze(0), - mode=self.mode.value if mode is None else GridSampleMode(mode).value, - padding_mode=self.padding_mode.value if padding_mode is None else GridSamplePadMode(padding_mode).value, + mode=self.mode if mode is None else GridSampleMode(mode), + padding_mode=self.padding_mode if padding_mode is None else GridSamplePadMode(padding_mode), align_corners=True, )[0] out_val, *_ = convert_to_dst_type(out, dst=img, dtype=np.float32) return out_val -class Affine(Transform): +class Affine(InvertibleTransform): """ Transform ``img`` given the affine parameters. A tutorial is available: https://github.com/Project-MONAI/tutorials/blob/0.6.0/modules/transforms_demo_2d.ipynb. @@ -1687,6 +2071,7 @@ class Affine(Transform): backend = list(set(AffineGrid.backend) & set(Resample.backend)) @deprecated_arg(name="as_tensor_output", since="0.6") + @deprecated_arg(name="norm_coords", since="0.8") def __init__( self, rotate_params: Optional[Union[Sequence[float], float]] = None, @@ -1695,8 +2080,9 @@ def __init__( scale_params: Optional[Union[Sequence[float], float]] = None, affine: Optional[NdarrayOrTensor] = None, spatial_size: Optional[Union[Sequence[int], int]] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.REFLECTION, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.REFLECTION, + normalized: bool = False, norm_coords: bool = True, as_tensor_output: bool = True, device: Optional[torch.device] = None, @@ -1741,11 +2127,11 @@ def __init__( padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``} Padding mode for outside grid values. Defaults to ``"reflection"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html - norm_coords: whether to normalize the coordinates from `[-(size-1)/2, (size-1)/2]` to - `[0, size - 1]` or `[-1, 1]` to be compatible with the underlying resampling API. - If the coordinates are generated by ``monai.transforms.utils.create_grid`` - and the ``affine`` doesn't include the normalization, this argument should be set to ``True``. - If the output `self.affine_grid` is already normalized, this argument should be set to ``False``. + 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. + 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 ``np.float32``. If ``None``, use the data type of input data. To be compatible with other modules, @@ -1754,6 +2140,9 @@ def __init__( .. deprecated:: 0.6.0 ``as_tensor_output`` is deprecated. + .. deprecated:: 0.8.1 + ``norm_coords`` is deprecated, please use ``normalized`` instead + (the new flag is a negation, i.e., ``norm_coords == not normalized``). """ self.affine_grid = AffineGrid( @@ -1766,18 +2155,19 @@ def __init__( device=device, ) self.image_only = image_only - self.resampler = Resample(norm_coords=norm_coords, device=device, dtype=dtype) + self.norm_coord = not normalized + self.resampler = Resample(norm_coords=self.norm_coord, device=device, dtype=dtype) self.spatial_size = spatial_size - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, spatial_size: Optional[Union[Sequence[int], int]] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, - ) -> Union[NdarrayOrTensor, Tuple[NdarrayOrTensor, NdarrayOrTensor]]: + mode: Optional[str] = None, + padding_mode: Optional[str] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, NdarrayOrTensor]]: """ Args: img: shape must be (num_channels, H, W[, D]), @@ -1796,14 +2186,60 @@ def __call__( Padding mode for outside grid values. Defaults to ``self.padding_mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html """ - sp_size = fall_back_tuple(self.spatial_size if spatial_size is None else spatial_size, img.shape[1:]) + img = convert_to_tensor(img, track_meta=get_track_meta()) + img_size = img.shape[1:] + sp_size = fall_back_tuple(self.spatial_size if spatial_size is None else spatial_size, img_size) + _mode = mode or self.mode + _padding_mode = padding_mode or self.padding_mode grid, affine = self.affine_grid(spatial_size=sp_size) - ret = self.resampler(img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode) - - return ret if self.image_only else (ret, affine) - - -class RandAffine(RandomizableTransform): + out = self.resampler(img, grid=grid, mode=_mode, padding_mode=_padding_mode) + if not isinstance(out, MetaTensor): + return out if self.image_only else (out, affine) + if not self.norm_coord: + warnings.warn("customized transform may not work with the metadata operation.") + if get_track_meta(): + out.meta = img.meta # type: ignore + self.update_meta(out, affine, img_size, sp_size) + self.push_transform( + out, orig_size=img_size, extra_info={"affine": affine, "mode": _mode, "padding_mode": _padding_mode} + ) + return out if self.image_only else (out, affine) + + @classmethod + def compute_w_affine(cls, affine, mat, img_size, sp_size): + r = len(affine) - 1 + mat = to_affine_nd(r, mat) + shift_1 = create_translate(r, [float(d - 1) / 2 for d in img_size[:r]]) + shift_2 = create_translate(r, [-float(d - 1) / 2 for d in sp_size[:r]]) + mat = shift_1 @ convert_data_type(mat, np.ndarray)[0] @ shift_2 + return affine @ convert_to_dst_type(mat, affine)[0] + + def update_meta(self, img, mat, img_size, sp_size): + affine = convert_data_type(img.affine, torch.Tensor)[0] + img.affine = Affine.compute_w_affine(affine, mat, img_size, sp_size) + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + orig_size = transform[TraceKeys.ORIG_SIZE] + # Create inverse transform + fwd_affine = transform[TraceKeys.EXTRA_INFO]["affine"] + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] + inv_affine = linalg_inv(fwd_affine) + inv_affine = convert_to_dst_type(inv_affine, data, dtype=inv_affine.dtype)[0] + + affine_grid = AffineGrid(affine=inv_affine) + grid, _ = affine_grid(orig_size) + # Apply inverse transform + out = self.resampler(data, grid, mode, padding_mode) + if not isinstance(out, MetaTensor): + out = MetaTensor(out) + out.meta = data.meta # type: ignore + self.update_meta(out, inv_affine, data.shape[1:], orig_size) + return out # type: ignore + + +class RandAffine(RandomizableTransform, InvertibleTransform): """ Random affine transform. A tutorial is available: https://github.com/Project-MONAI/tutorials/blob/0.6.0/modules/transforms_demo_2d.ipynb. @@ -1821,8 +2257,8 @@ def __init__( translate_range: RandRange = None, scale_range: RandRange = None, spatial_size: Optional[Union[Sequence[int], int]] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.REFLECTION, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.REFLECTION, cache_grid: bool = False, as_tensor_output: bool = True, device: Optional[torch.device] = None, @@ -1892,8 +2328,8 @@ def __init__( self.spatial_size = spatial_size self.cache_grid = cache_grid self._cached_grid = self._init_identity_cache() - self.mode: GridSampleMode = GridSampleMode(mode) - self.padding_mode: GridSamplePadMode = GridSamplePadMode(padding_mode) + self.mode: str = GridSampleMode(mode) + self.padding_mode: str = GridSamplePadMode(padding_mode) def _init_identity_cache(self): """ @@ -1950,12 +2386,13 @@ def randomize(self, data: Optional[Any] = None) -> None: def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, spatial_size: Optional[Union[Sequence[int], int]] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, randomize: bool = True, - ) -> NdarrayOrTensor: + grid=None, + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W[, D]), @@ -1971,25 +2408,70 @@ def __call__( Padding mode for outside grid values. Defaults to ``self.padding_mode``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html randomize: whether to execute `randomize()` function first, default to True. + grid: precomputed grid to be used (mainly to accelerate `RandAffined`). """ if randomize: self.randomize() - # if not doing transform and spatial size doesn't change, nothing to do # except convert to float and device sp_size = fall_back_tuple(self.spatial_size if spatial_size is None else spatial_size, img.shape[1:]) do_resampling = self._do_transform or (sp_size != ensure_tuple(img.shape[1:])) + _mode = mode or self.mode + _padding_mode = padding_mode or self.padding_mode + img = convert_to_tensor(img, track_meta=get_track_meta()) if not do_resampling: - img, *_ = convert_data_type(img, dtype=torch.float32, device=self.resampler.device) - grid = self.get_identity_grid(sp_size) - if self._do_transform: - grid = self.rand_affine_grid(grid=grid) - out: NdarrayOrTensor = self.resampler( - img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode - ) + out: torch.Tensor = convert_data_type(img, dtype=torch.float32, device=self.resampler.device)[0] + else: + if grid is None: + grid = self.get_identity_grid(sp_size) + if self._do_transform: + grid = self.rand_affine_grid(grid=grid, randomize=randomize) + out = self.resampler(img=img, grid=grid, mode=_mode, padding_mode=_padding_mode) + mat = self.rand_affine_grid.get_transformation_matrix() + out = convert_to_tensor(out, track_meta=get_track_meta()) + if get_track_meta(): + self.push_transform( + out, + orig_size=img.shape[1:], + extra_info={ + "affine": mat, + "mode": _mode, + "padding_mode": _padding_mode, + "do_resampling": do_resampling, + }, + ) + self.update_meta(out, mat, img.shape[1:], sp_size) return out + def update_meta(self, img, mat, img_size, sp_size): + affine = convert_data_type(img.affine, torch.Tensor)[0] + img.affine = Affine.compute_w_affine(affine, mat, img_size, sp_size) + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + transform = self.pop_transform(data) + # if transform was not performed nothing to do. + if not transform[TraceKeys.EXTRA_INFO]["do_resampling"]: + return data + orig_size = transform[TraceKeys.ORIG_SIZE] + orig_size = fall_back_tuple(orig_size, data.shape[1:]) + # Create inverse transform + fwd_affine = transform[TraceKeys.EXTRA_INFO]["affine"] + mode = transform[TraceKeys.EXTRA_INFO]["mode"] + padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] + inv_affine = linalg_inv(fwd_affine) + inv_affine = convert_to_dst_type(inv_affine, data, dtype=inv_affine.dtype)[0] + affine_grid = AffineGrid(affine=inv_affine) + grid, _ = affine_grid(orig_size) + + # Apply inverse transform + out = self.resampler(data, grid, mode, padding_mode) + if not isinstance(out, MetaTensor): + out = MetaTensor(out) + out.meta = data.meta # type: ignore + self.update_meta(out, inv_affine, data.shape[1:], orig_size) + return out # type: ignore + class Rand2DElastic(RandomizableTransform): """ @@ -2011,8 +2493,8 @@ def __init__( translate_range: RandRange = None, scale_range: RandRange = None, spatial_size: Optional[Union[Tuple[int, int], int]] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.REFLECTION, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.REFLECTION, as_tensor_output: bool = False, device: Optional[torch.device] = None, ) -> None: @@ -2066,9 +2548,7 @@ def __init__( """ RandomizableTransform.__init__(self, prob) - self.deform_grid = RandDeformGrid( - spacing=spacing, magnitude_range=magnitude_range, as_tensor_output=True, device=device - ) + self.deform_grid = RandDeformGrid(spacing=spacing, magnitude_range=magnitude_range, device=device) self.rand_affine_grid = RandAffineGrid( rotate_range=rotate_range, shear_range=shear_range, @@ -2080,8 +2560,8 @@ def __init__( self.device = device self.spatial_size = spatial_size - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) def set_random_state( self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None @@ -2100,12 +2580,12 @@ def randomize(self, spatial_size: Sequence[int]) -> None: def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, spatial_size: Optional[Union[Tuple[int, int], int]] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, randomize: bool = True, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W), @@ -2127,7 +2607,7 @@ def __call__( if self._do_transform: grid = self.deform_grid(spatial_size=sp_size) grid = self.rand_affine_grid(grid=grid) - grid = torch.nn.functional.interpolate( # type: ignore + grid = torch.nn.functional.interpolate( recompute_scale_factor=True, input=grid.unsqueeze(0), scale_factor=list(ensure_tuple(self.deform_grid.spacing)), @@ -2138,7 +2618,7 @@ def __call__( else: _device = img.device if isinstance(img, torch.Tensor) else self.device grid = create_grid(spatial_size=sp_size, device=_device, backend="torch") - out: NdarrayOrTensor = self.resampler( + out: torch.Tensor = self.resampler( img, grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode ) return out @@ -2164,8 +2644,8 @@ def __init__( translate_range: RandRange = None, scale_range: RandRange = None, spatial_size: Optional[Union[Tuple[int, int, int], int]] = None, - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.REFLECTION, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.REFLECTION, as_tensor_output: bool = False, device: Optional[torch.device] = None, ) -> None: @@ -2234,8 +2714,8 @@ def __init__( self.sigma_range = sigma_range self.magnitude_range = magnitude_range self.spatial_size = spatial_size - self.mode: GridSampleMode = look_up_option(mode, GridSampleMode) - self.padding_mode: GridSamplePadMode = look_up_option(padding_mode, GridSamplePadMode) + self.mode: str = look_up_option(mode, GridSampleMode) + self.padding_mode: str = look_up_option(padding_mode, GridSamplePadMode) self.device = device self.rand_offset: np.ndarray @@ -2260,12 +2740,12 @@ def randomize(self, grid_size: Sequence[int]) -> None: def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, spatial_size: Optional[Union[Tuple[int, int, int], int]] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, + mode: Optional[str] = None, + padding_mode: Optional[str] = None, randomize: bool = True, - ) -> NdarrayOrTensor: + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W, D), @@ -2293,7 +2773,7 @@ def __call__( offset = torch.as_tensor(self.rand_offset, device=_device).unsqueeze(0) grid[:3] += gaussian(offset)[0] * self.magnitude grid = self.rand_affine_grid(grid=grid) - out: NdarrayOrTensor = self.resampler( + out: torch.Tensor = self.resampler( img, grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode ) return out @@ -2307,8 +2787,8 @@ def __init__( self, num_cells: Union[Tuple[int], int], distort_steps: Sequence[Sequence[float]], - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, device: Optional[torch.device] = None, ) -> None: """ @@ -2336,11 +2816,11 @@ def __init__( def __call__( self, - img: NdarrayOrTensor, + img: torch.Tensor, distort_steps: Optional[Sequence[Sequence]] = None, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, - ) -> NdarrayOrTensor: + mode: Optional[str] = None, + padding_mode: Optional[str] = None, + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W[, D]). @@ -2394,8 +2874,8 @@ def __init__( num_cells: Union[Tuple[int], int] = 5, prob: float = 0.1, distort_limit: Union[Tuple[float, float], float] = (-0.03, 0.03), - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, device: Optional[torch.device] = None, ) -> None: """ @@ -2438,12 +2918,8 @@ def randomize(self, spatial_shape: Sequence[int]) -> None: ) def __call__( - self, - img: NdarrayOrTensor, - mode: Optional[Union[GridSampleMode, str]] = None, - padding_mode: Optional[Union[GridSamplePadMode, str]] = None, - randomize: bool = True, - ) -> NdarrayOrTensor: + self, img: torch.Tensor, mode: Optional[str] = None, padding_mode: Optional[str] = None, randomize: bool = True + ) -> torch.Tensor: """ Args: img: shape must be (num_channels, H, W[, D]). @@ -2458,5 +2934,281 @@ def __call__( if randomize: self.randomize(img.shape[1:]) if not self._do_transform: - return img + return convert_to_tensor(img, track_meta=get_track_meta()) # type: ignore return self.grid_distortion(img, distort_steps=self.distort_steps, mode=mode, padding_mode=padding_mode) + + +class GridSplit(Transform): + """ + Split the image into patches based on the provided grid in 2D. + + Args: + grid: a tuple define the shape of the grid upon which the image is split. Defaults to (2, 2) + size: a tuple or an integer that defines the output patch sizes. + If it's an integer, the value will be repeated for each dimension. + The default is None, where the patch size will be inferred from the grid shape. + + Example: + Given an image (torch.Tensor or numpy.ndarray) with size of (3, 10, 10) and a grid of (2, 2), + it will return a Tensor or array with the size of (4, 3, 5, 5). + Here, if the `size` is provided, the returned shape will be (4, 3, size, size) + + Note: This transform currently support only image with two spatial dimensions. + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, grid: Tuple[int, int] = (2, 2), size: Optional[Union[int, Tuple[int, int]]] = None): + # Grid size + self.grid = grid + + # Patch size + self.size = None if size is None else ensure_tuple_rep(size, len(self.grid)) + + def __call__( + self, image: NdarrayOrTensor, size: Optional[Union[int, Tuple[int, int], np.ndarray]] = None + ) -> List[NdarrayOrTensor]: + input_size = self.size if size is None else ensure_tuple_rep(size, len(self.grid)) + + if self.grid == (1, 1) and input_size is None: + return [image] + + split_size, steps = self._get_params(image.shape[1:], input_size) + patches: List[NdarrayOrTensor] + if isinstance(image, torch.Tensor): + unfolded_image = ( + image.unfold(1, split_size[0], steps[0]) + .unfold(2, split_size[1], steps[1]) + .flatten(1, 2) + .transpose(0, 1) + ) + # Make a list of contiguous patches + patches = [p.contiguous() for p in unfolded_image] + elif isinstance(image, np.ndarray): + x_step, y_step = steps + c_stride, x_stride, y_stride = image.strides + n_channels = image.shape[0] + strided_image = as_strided( + image, + shape=(*self.grid, n_channels, split_size[0], split_size[1]), + strides=(x_stride * x_step, y_stride * y_step, c_stride, x_stride, y_stride), + ) + # Flatten the first two dimensions + strided_image = strided_image.reshape(-1, *strided_image.shape[2:]) + # Make a list of contiguous patches + patches = [np.ascontiguousarray(p) for p in strided_image] + else: + raise ValueError(f"Input type [{type(image)}] is not supported.") + + return patches + + def _get_params( + self, image_size: Union[Sequence[int], np.ndarray], size: Optional[Union[Sequence[int], np.ndarray]] = None + ): + """ + Calculate the size and step required for splitting the image + Args: + The size of the input image + """ + if size is None: + # infer each sub-image size from the image size and the grid + size = tuple(image_size[i] // self.grid[i] for i in range(len(self.grid))) + + if any(size[i] > image_size[i] for i in range(len(self.grid))): + raise ValueError(f"The image size ({image_size})is smaller than the requested split size ({size})") + + steps = tuple( + (image_size[i] - size[i]) // (self.grid[i] - 1) if self.grid[i] > 1 else image_size[i] + for i in range(len(self.grid)) + ) + + return size, steps + + +class GridPatch(Transform): + """ + Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps. + It can sort the patches and return all or a subset of them. + + Args: + patch_size: size of patches to generate slices for, 0 or None selects whole dimension + offset: offset of starting position in the array, default is 0 for each dimension. + num_patches: number of patches to return. Defaults to None, which returns all the available patches. + If the required patches are more than the available patches, padding will be applied. + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`), + lowest values (`"min"`), or in their default order (`None`). Default to None. + threshold: a value to keep only the patches whose sum of intensities are less than the threshold. + Defaults to no filtering. + pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``. + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + patch_size: Sequence[int], + offset: Optional[Sequence[int]] = None, + num_patches: Optional[int] = None, + overlap: Union[Sequence[float], float] = 0.0, + sort_fn: Optional[str] = None, + threshold: Optional[float] = None, + pad_mode: str = PytorchPadMode.CONSTANT, + **pad_kwargs, + ): + self.patch_size = ensure_tuple(patch_size) + self.offset = ensure_tuple(offset) if offset else (0,) * len(self.patch_size) + self.pad_mode: Optional[NumpyPadMode] = convert_pad_mode(dst=np.zeros(1), mode=pad_mode) if pad_mode else None + self.pad_kwargs = pad_kwargs + self.overlap = overlap + self.num_patches = num_patches + self.sort_fn = sort_fn.lower() if sort_fn else None + self.threshold = threshold + + def filter_threshold(self, image_np: np.ndarray, locations: np.ndarray): + """ + Filter the patches and their locations according to a threshold + Args: + image_np: a numpy.ndarray representing a stack of patches + locations: a numpy.ndarray representing the stack of location of each patch + """ + if self.threshold is not None: + n_dims = len(image_np.shape) + idx = np.argwhere(image_np.sum(axis=tuple(range(1, n_dims))) < self.threshold).reshape(-1) + image_np = image_np[idx] + locations = locations[idx] + return image_np, locations + + def filter_count(self, image_np: np.ndarray, locations: np.ndarray): + """ + Sort the patches based on the sum of their intensity, and just keep `self.num_patches` of them. + Args: + image_np: a numpy.ndarray representing a stack of patches + locations: a numpy.ndarray representing the stack of location of each patch + """ + if self.sort_fn is None: + image_np = image_np[: self.num_patches] + locations = locations[: self.num_patches] + elif self.num_patches is not None: + n_dims = len(image_np.shape) + if self.sort_fn == GridPatchSort.MIN: + idx = np.argsort(image_np.sum(axis=tuple(range(1, n_dims)))) + elif self.sort_fn == GridPatchSort.MAX: + idx = np.argsort(-image_np.sum(axis=tuple(range(1, n_dims)))) + else: + raise ValueError(f'`sort_fn` should be either "min", "max" or None! {self.sort_fn} provided!') + idx = idx[: self.num_patches] + image_np = image_np[idx] + locations = locations[idx] + return image_np, locations + + def __call__(self, array: NdarrayOrTensor): + # create the patch iterator which sweeps the image row-by-row + array_np, *_ = convert_data_type(array, np.ndarray) + patch_iterator = iter_patch( + array_np, + patch_size=(None,) + self.patch_size, # expand to have the channel dim + start_pos=(0,) + self.offset, # expand to have the channel dim + overlap=self.overlap, + copy_back=False, + mode=self.pad_mode, + **self.pad_kwargs, + ) + patches = list(zip(*patch_iterator)) + patched_image = np.array(patches[0]) + locations = np.array(patches[1])[:, 1:, 0] # only keep the starting location + + # Filter patches + if self.num_patches: + patched_image, locations = self.filter_count(patched_image, locations) + elif self.threshold: + patched_image, locations = self.filter_threshold(patched_image, locations) + + # Convert to original data type + output = list( + zip( + convert_to_dst_type(src=patched_image, dst=array)[0], + convert_to_dst_type(src=locations, dst=array, dtype=int)[0], + ) + ) + + # Pad the patch list to have the requested number of patches + if self.num_patches and len(output) < self.num_patches: + patch = convert_to_dst_type( + src=np.full((array.shape[0], *self.patch_size), self.pad_kwargs.get("constant_values", 0)), dst=array + )[0] + start_location = convert_to_dst_type(src=np.zeros(len(self.patch_size)), dst=array)[0] + output += [(patch, start_location)] * (self.num_patches - len(output)) + + return output + + +class RandGridPatch(GridPatch, RandomizableTransform): + """ + Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps, + and with random offset for the minimal corner of the image, (0,0) for 2D and (0,0,0) for 3D. + It can sort the patches and return all or a subset of them. + + Args: + patch_size: size of patches to generate slices for, 0 or None selects whole dimension + min_offset: the minimum range of offset to be selected randomly. Defaults to 0. + max_offset: the maximum range of offset to be selected randomly. + Defaults to image size modulo patch size. + num_patches: number of patches to return. Defaults to None, which returns all the available patches. + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`), + lowest values (`"min"`), or in their default order (`None`). Default to None. + threshold: a value to keep only the patches whose sum of intensities are less than the threshold. + Defaults to no filtering. + pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``. + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + patch_size: Sequence[int], + min_offset: Optional[Union[Sequence[int], int]] = None, + max_offset: Optional[Union[Sequence[int], int]] = None, + num_patches: Optional[int] = None, + overlap: Union[Sequence[float], float] = 0.0, + sort_fn: Optional[str] = None, + threshold: Optional[float] = None, + pad_mode: str = PytorchPadMode.CONSTANT, + **pad_kwargs, + ): + super().__init__( + patch_size=patch_size, + offset=(), + num_patches=num_patches, + overlap=overlap, + sort_fn=sort_fn, + threshold=threshold, + pad_mode=pad_mode, + **pad_kwargs, + ) + self.min_offset = min_offset + self.max_offset = max_offset + + def randomize(self, array): + if self.min_offset is None: + min_offset = (0,) * len(self.patch_size) + else: + min_offset = ensure_tuple_rep(self.min_offset, len(self.patch_size)) + if self.max_offset is None: + max_offset = tuple(s % p for s, p in zip(array.shape[1:], self.patch_size)) + else: + max_offset = ensure_tuple_rep(self.max_offset, len(self.patch_size)) + + self.offset = tuple(self.R.randint(low=low, high=high + 1) for low, high in zip(min_offset, max_offset)) + + def __call__(self, array: NdarrayOrTensor, randomize: bool = True): + if randomize: + self.randomize(array) + return super().__call__(array) diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py index d42a11fd2f..c809d38ba0 100644 --- a/monai/transforms/spatial/dictionary.py +++ b/monai/transforms/spatial/dictionary.py @@ -16,31 +16,31 @@ """ from copy import deepcopy -from enum import Enum from typing import Any, Dict, Hashable, List, Mapping, Optional, Sequence, Tuple, Union import numpy as np import torch -from monai.config import DtypeLike, KeysCollection +from monai.config import DtypeLike, KeysCollection, SequenceStr from monai.config.type_definitions import NdarrayOrTensor -from monai.data.utils import affine_to_spacing -from monai.networks.layers import AffineTransform +from monai.data.meta_obj import get_track_meta +from monai.data.meta_tensor import MetaTensor from monai.networks.layers.simplelayers import GaussianFilter -from monai.transforms.croppad.array import CenterSpatialCrop, SpatialPad +from monai.transforms.croppad.array import CenterSpatialCrop from monai.transforms.inverse import InvertibleTransform from monai.transforms.spatial.array import ( Affine, - AffineGrid, Flip, GridDistortion, + GridPatch, + GridSplit, Orientation, Rand2DElastic, Rand3DElastic, RandAffine, RandAxisFlip, - RandFlip, RandGridDistortion, + RandGridPatch, RandRotate, RandZoom, ResampleToMatch, @@ -58,15 +58,16 @@ GridSamplePadMode, InterpolateMode, NumpyPadMode, - PytorchPadMode, + WSIPatchKeys, + convert_to_tensor, ensure_tuple, ensure_tuple_rep, fall_back_tuple, + first, ) from monai.utils.deprecate_utils import deprecated_arg -from monai.utils.enums import PostFix, TraceKeys +from monai.utils.enums import PytorchPadMode, TraceKeys from monai.utils.module import optional_import -from monai.utils.type_conversion import convert_data_type, convert_to_dst_type nib, _ = optional_import("nibabel") @@ -129,14 +130,17 @@ "ZoomDict", "RandZoomD", "RandZoomDict", + "GridSplitd", + "GridSplitD", + "GridSplitDict", + "GridPatchd", + "GridPatchD", + "GridPatchDict", + "RandGridPatchd", + "RandGridPatchD", + "RandGridPatchDict", ] -GridSampleModeSequence = Union[Sequence[Union[GridSampleMode, str]], GridSampleMode, str] -GridSamplePadModeSequence = Union[Sequence[Union[GridSamplePadMode, str]], GridSamplePadMode, str] -InterpolateModeSequence = Union[Sequence[Union[InterpolateMode, str]], InterpolateMode, str] -PadModeSequence = Union[Sequence[Union[NumpyPadMode, PytorchPadMode, str]], NumpyPadMode, PytorchPadMode, str] -DEFAULT_POST_FIX = PostFix.meta() - class SpatialResampled(MapTransform, InvertibleTransform): """ @@ -155,17 +159,20 @@ class SpatialResampled(MapTransform, InvertibleTransform): backend = SpatialResample.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") + @deprecated_arg(name="meta_src_keys", since="0.9") def __init__( self, keys: KeysCollection, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.BORDER, align_corners: Union[Sequence[bool], bool] = False, dtype: Union[Sequence[DtypeLike], DtypeLike] = np.float64, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", meta_src_keys: Optional[KeysCollection] = "src_affine", - meta_dst_keys: Optional[KeysCollection] = "dst_affine", + dst_keys: Optional[KeysCollection] = "dst_affine", allow_missing_keys: bool = False, ) -> None: """ @@ -186,17 +193,7 @@ def __init__( If None, use the data type of input data. To be compatible with other modules, the output data type is always ``np.float32``. It also can be a sequence of dtypes, each element corresponds to a key in ``keys``. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, affine, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys=None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - For example, to handle key `image`, read/write affine matrices from the - metadata `image_meta_dict` dictionary's `affine` field. - meta_src_keys: the key of the corresponding ``src_affine`` in the metadata dictionary. - meta_dst_keys: the key of the corresponding ``dst_affine`` in the metadata dictionary. + dst_keys: the key of the corresponding ``dst_affine`` in the metadata dictionary. allow_missing_keys: don't raise exception if key is missing. """ super().__init__(keys, allow_missing_keys) @@ -205,90 +202,28 @@ def __init__( self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.dtype = ensure_tuple_rep(dtype, len(self.keys)) - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - self.meta_src_keys = ensure_tuple_rep(meta_src_keys, len(self.keys)) - self.meta_dst_keys = ensure_tuple_rep(meta_dst_keys, len(self.keys)) - - def __call__( - self, data: Mapping[Union[Hashable, str], Dict[str, NdarrayOrTensor]] - ) -> Dict[Hashable, NdarrayOrTensor]: + self.dst_keys = ensure_tuple_rep(dst_keys, len(self.keys)) + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d: Dict = dict(data) - for (key, mode, padding_mode, align_corners, dtype, *metakeyinfo) in self.key_iterator( - d, - self.mode, - self.padding_mode, - self.align_corners, - self.dtype, - self.meta_keys, - self.meta_key_postfix, - self.meta_src_keys, - self.meta_dst_keys, + for (key, mode, padding_mode, align_corners, dtype, dst_key) in self.key_iterator( + d, self.mode, self.padding_mode, self.align_corners, self.dtype, self.dst_keys ): - meta_key, meta_key_postfix, meta_src_key, meta_dst_key = metakeyinfo - meta_key = meta_key or f"{key}_{meta_key_postfix}" - # create metadata if necessary - if meta_key not in d: - d[meta_key] = {meta_src_key: None, meta_dst_key: None} - meta_data = d[meta_key] - original_spatial_shape = d[key].shape[1:] - d[key], meta_data[meta_dst_key] = self.sp_transform( # write dst affine because the dtype might change + d[key] = self.sp_transform( img=d[key], - src_affine=meta_data[meta_src_key], - dst_affine=meta_data[meta_dst_key], + dst_affine=d[dst_key], spatial_size=None, # None means shape auto inferred mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype, ) - meta_data[meta_dst_key], meta_data[meta_src_key] = meta_data[meta_src_key], meta_data[meta_dst_key] - self.push_transform( - d, - key, - extra_info={ - "meta_key": meta_key, - "meta_src_key": meta_src_key, - "meta_dst_key": meta_dst_key, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - orig_size=original_spatial_shape, - ) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - for key, dtype in self.key_iterator(d, self.dtype): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - meta_data = d[transform[TraceKeys.EXTRA_INFO]["meta_key"]] - src_key = transform[TraceKeys.EXTRA_INFO]["meta_src_key"] - dst_key = transform[TraceKeys.EXTRA_INFO]["meta_dst_key"] - src_affine = meta_data[src_key] - dst_affine = meta_data[dst_key] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - orig_size = transform[TraceKeys.ORIG_SIZE] - inverse_transform = SpatialResample() - # Apply inverse - d[key], dst_affine = inverse_transform( - img=d[key], - src_affine=src_affine, - dst_affine=dst_affine, - mode=mode, - padding_mode=padding_mode, - align_corners=False if align_corners == TraceKeys.NONE else align_corners, - dtype=dtype, - spatial_size=orig_size, - ) - meta_data[src_key], meta_data[dst_key] = dst_affine, meta_data[src_key] # type: ignore - # Remove the applied transform - self.pop_transform(d, key) + for key in self.key_iterator(d): + d[key] = self.sp_transform.inverse(d[key]) return d @@ -297,12 +232,14 @@ class ResampleToMatchd(MapTransform, InvertibleTransform): backend = ResampleToMatch.backend + @deprecated_arg(name="template_key", since="0.9") def __init__( self, keys: KeysCollection, - template_key: str, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER, + key_dst: str, + template_key: Optional[str] = None, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.BORDER, align_corners: Union[Sequence[bool], bool] = False, dtype: Union[Sequence[DtypeLike], DtypeLike] = np.float64, allow_missing_keys: bool = False, @@ -310,7 +247,7 @@ def __init__( """ Args: keys: keys of the corresponding items to be transformed. - template_key: key to meta data that output should be resampled to match. + key_dst: key of image to resample to match. mode: {``"bilinear"``, ``"nearest"``} Interpolation mode to calculate output values. Defaults to ``"bilinear"``. See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample @@ -329,79 +266,32 @@ def __init__( allow_missing_keys: don't raise exception if key is missing. """ super().__init__(keys, allow_missing_keys) - self.template_key = template_key + self.key_dst = key_dst self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.dtype = ensure_tuple_rep(dtype, len(self.keys)) self.resampler = ResampleToMatch() - def __call__(self, data): + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - dst_meta = d[self.template_key] for (key, mode, padding_mode, align_corners, dtype) in self.key_iterator( d, self.mode, self.padding_mode, self.align_corners, self.dtype ): - src_meta_key = PostFix.meta(key) - src_meta = d[src_meta_key] - - orig_spatial_shape = d[key].shape[1:] - orig_meta = deepcopy(src_meta) - - img, new_meta = self.resampler( + d[key] = self.resampler( img=d[key], - src_meta=src_meta, - dst_meta=dst_meta, + img_dst=d[self.key_dst], mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype, ) - d[key] = img - d[src_meta_key] = new_meta - - # track the transform for the inverse - self.push_transform( - d, - key, - extra_info={ - "orig_meta": orig_meta, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - orig_size=orig_spatial_shape, - ) - return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - for key, dtype in self.key_iterator(d, self.dtype): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - orig_meta = transform[TraceKeys.EXTRA_INFO]["orig_meta"] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - - src_meta_key = PostFix.meta(key) - src_meta = d[src_meta_key] - - img, new_meta = self.resampler( - img=d[key], - src_meta=src_meta, # type: ignore - dst_meta=orig_meta, - mode=mode, - padding_mode=padding_mode, - align_corners=align_corners, - dtype=dtype, - ) - d[key] = img - d[src_meta_key] = new_meta - - # Remove the applied transform - self.pop_transform(d, key) + for key in self.key_iterator(d): + d[key] = self.resampler.inverse(d[key]) return d @@ -421,17 +311,19 @@ class Spacingd(MapTransform, InvertibleTransform): backend = Spacing.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") def __init__( self, keys: KeysCollection, pixdim: Union[Sequence[float], float], diagonal: bool = False, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.BORDER, align_corners: Union[Sequence[bool], bool] = False, dtype: Union[Sequence[DtypeLike], DtypeLike] = np.float64, meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_missing_keys: bool = False, ) -> None: """ @@ -470,20 +362,8 @@ def __init__( If None, use the data type of input data. To be compatible with other modules, the output data type is always ``np.float32``. It also can be a sequence of dtypes, each element corresponds to a key in ``keys``. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, affine, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys=None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - For example, to handle key `image`, read/write affine matrices from the - metadata `image_meta_dict` dictionary's `affine` field. allow_missing_keys: don't raise exception if key is missing. - Raises: - TypeError: When ``meta_key_postfix`` is not a ``str``. - """ super().__init__(keys, allow_missing_keys) self.spacing_transform = Spacing(pixdim, diagonal=diagonal) @@ -491,82 +371,22 @@ def __init__( self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.dtype = ensure_tuple_rep(dtype, len(self.keys)) - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - - def __call__( - self, data: Mapping[Union[Hashable, str], Dict[str, NdarrayOrTensor]] - ) -> Dict[Hashable, NdarrayOrTensor]: + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d: Dict = dict(data) - for key, mode, padding_mode, align_corners, dtype, meta_key, meta_key_postfix in self.key_iterator( - d, self.mode, self.padding_mode, self.align_corners, self.dtype, self.meta_keys, self.meta_key_postfix + for key, mode, padding_mode, align_corners, dtype in self.key_iterator( + d, self.mode, self.padding_mode, self.align_corners, self.dtype ): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - # create metadata if necessary - if meta_key not in d: - d[meta_key] = {"affine": None} - meta_data = d[meta_key] # resample array of each corresponding key - # using affine fetched from d[affine_key] - original_spatial_shape = d[key].shape[1:] - d[key], old_affine, new_affine = self.spacing_transform( - data_array=d[key], - affine=meta_data["affine"], - mode=mode, - padding_mode=padding_mode, - align_corners=align_corners, - dtype=dtype, - ) - self.push_transform( - d, - key, - extra_info={ - "meta_key": meta_key, - "old_affine": old_affine, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - orig_size=original_spatial_shape, + d[key] = self.spacing_transform( + data_array=d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype ) - # set the 'affine' key - meta_data["affine"] = new_affine return d def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: d = deepcopy(dict(data)) - for key, dtype in self.key_iterator(d, self.dtype): - transform = self.get_most_recent_transform(d, key) - if self.spacing_transform.diagonal: - raise RuntimeError( - "Spacingd:inverse not yet implemented for diagonal=True. " - + "Please raise a github issue if you need this feature" - ) - # Create inverse transform - meta_data = d[transform[TraceKeys.EXTRA_INFO]["meta_key"]] - old_affine = np.array(transform[TraceKeys.EXTRA_INFO]["old_affine"]) - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - orig_size = transform[TraceKeys.ORIG_SIZE] - orig_pixdim = affine_to_spacing(old_affine, -1) - inverse_transform = Spacing(orig_pixdim, diagonal=self.spacing_transform.diagonal) - # Apply inverse - d[key], _, new_affine = inverse_transform( - data_array=d[key], - affine=meta_data["affine"], # type: ignore - mode=mode, - padding_mode=padding_mode, - align_corners=False if align_corners == TraceKeys.NONE else align_corners, - dtype=dtype, - output_spatial_shape=orig_size, - ) - meta_data["affine"] = new_affine # type: ignore - # Remove the applied transform - self.pop_transform(d, key) - + for key in self.key_iterator(d): + d[key] = self.spacing_transform.inverse(d[key]) return d @@ -574,12 +394,6 @@ class Orientationd(MapTransform, InvertibleTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.Orientation`. - This transform assumes the ``data`` dictionary has a key for the input - data's metadata and contains `affine` field. The key is formed by ``key_{meta_key_postfix}``. - - After reorienting the input array, this transform will write the new affine - to the `affine` field of metadata which is formed by ``key_{meta_key_postfix}``. - This transform assumes the channel-first input format. In the case of using this transform for normalizing the orientations of images, it should be used before any anisotropic spatial transforms. @@ -587,6 +401,8 @@ class Orientationd(MapTransform, InvertibleTransform): backend = Orientation.backend + @deprecated_arg(name="meta_keys", since="0.9") + @deprecated_arg(name="meta_key_postfix", since="0.9") def __init__( self, keys: KeysCollection, @@ -594,7 +410,7 @@ def __init__( as_closest_canonical: bool = False, labels: Optional[Sequence[Tuple[str, str]]] = (("L", "R"), ("P", "A"), ("I", "S")), meta_keys: Optional[KeysCollection] = None, - meta_key_postfix: str = DEFAULT_POST_FIX, + meta_key_postfix: str = "meta_dict", allow_missing_keys: bool = False, ) -> None: """ @@ -608,65 +424,25 @@ def __init__( labels: optional, None or sequence of (2,) sequences (2,) sequences are labels for (beginning, end) of output axis. Defaults to ``(('L', 'R'), ('P', 'A'), ('I', 'S'))``. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, affine, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. - For example, to handle key `image`, read/write affine matrices from the - metadata `image_meta_dict` dictionary's `affine` field. allow_missing_keys: don't raise exception if key is missing. - Raises: - TypeError: When ``meta_key_postfix`` is not a ``str``. - See Also: `nibabel.orientations.ornt2axcodes`. """ super().__init__(keys, allow_missing_keys) self.ornt_transform = Orientation(axcodes=axcodes, as_closest_canonical=as_closest_canonical, labels=labels) - if not isinstance(meta_key_postfix, str): - raise TypeError(f"meta_key_postfix must be a str but is {type(meta_key_postfix).__name__}.") - self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) - if len(self.keys) != len(self.meta_keys): - raise ValueError("meta_keys should have the same length as keys.") - self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - - def __call__( - self, data: Mapping[Union[Hashable, str], Dict[str, NdarrayOrTensor]] - ) -> Dict[Hashable, NdarrayOrTensor]: + + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d: Dict = dict(data) - for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): - meta_key = meta_key or f"{key}_{meta_key_postfix}" - # create metadata if necessary - if meta_key not in d: - d[meta_key] = {"affine": None} - meta_data = d[meta_key] - d[key], old_affine, new_affine = self.ornt_transform(d[key], affine=meta_data["affine"]) - self.push_transform(d, key, extra_info={"meta_key": meta_key, "old_affine": old_affine}) - d[meta_key]["affine"] = new_affine + for key in self.key_iterator(d): + d[key] = self.ornt_transform(d[key]) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - meta_data: Dict = d[transform[TraceKeys.EXTRA_INFO]["meta_key"]] # type: ignore - orig_affine = transform[TraceKeys.EXTRA_INFO]["old_affine"] - orig_axcodes = nib.orientations.aff2axcodes(orig_affine) - inverse_transform = Orientation( - axcodes=orig_axcodes, as_closest_canonical=False, labels=self.ornt_transform.labels - ) - # Apply inverse - d[key], _, new_affine = inverse_transform(d[key], affine=meta_data["affine"]) - meta_data["affine"] = new_affine - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.ornt_transform.inverse(d[key]) return d @@ -690,27 +466,16 @@ def __init__( super().__init__(keys, allow_missing_keys) self.rotator = Rotate90(k, spatial_axes) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key in self.key_iterator(d): - self.push_transform(d, key) d[key] = self.rotator(d[key]) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - _ = self.get_most_recent_transform(d, key) - # Create inverse transform - spatial_axes = self.rotator.spatial_axes - num_times_rotated = self.rotator.k - num_times_to_rotate = 4 - num_times_rotated - inverse_transform = Rotate90(num_times_to_rotate, spatial_axes) - # Apply inverse - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.rotator.inverse(d[key]) return d @@ -755,7 +520,7 @@ def randomize(self, data: Optional[Any] = None) -> None: self._rand_k = self.R.randint(self.max_k) + 1 super().randomize(None) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, torch.Tensor]: self.randomize() d = dict(data) @@ -763,26 +528,20 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable # to be compatible with the random status of some previous integration tests rotator = Rotate90(self._rand_k, self.spatial_axes) for key in self.key_iterator(d): - if self._do_transform: - d[key] = rotator(d[key]) - self.push_transform(d, key, extra_info={"rand_k": self._rand_k}) + d[key] = rotator(d[key]) if self._do_transform else convert_to_tensor(d[key], track_meta=get_track_meta()) + if get_track_meta(): + xform = self.pop_transform(d[key], check=False) if self._do_transform else {} + self.push_transform(d[key], extra_info=xform) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Check if random transform was actually performed (based on `prob`) - if transform[TraceKeys.DO_TRANSFORM]: - # Create inverse transform - num_times_rotated = transform[TraceKeys.EXTRA_INFO]["rand_k"] - num_times_to_rotate = 4 - num_times_rotated - inverse_transform = Rotate90(num_times_to_rotate, self.spatial_axes) - # Apply inverse - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + if not isinstance(d[key], MetaTensor): + continue + xform = self.pop_transform(d[key]) + if xform[TraceKeys.DO_TRANSFORM]: + d[key] = Rotate90().inverse_transform(d[key], xform[TraceKeys.EXTRA_INFO]) return d @@ -802,7 +561,7 @@ class Resized(MapTransform, InvertibleTransform): which must be an int number in this case, keeping the aspect ratio of the initial image, refer to: https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/ #albumentations.augmentations.geometric.resize.LongestMaxSize. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html It also can be a sequence of string, each element corresponds to a key in ``keys``. @@ -820,7 +579,7 @@ def __init__( keys: KeysCollection, spatial_size: Union[Sequence[int], int], size_mode: str = "all", - mode: InterpolateModeSequence = InterpolateMode.AREA, + mode: SequenceStr = InterpolateMode.AREA, align_corners: Union[Sequence[Optional[bool]], Optional[bool]] = None, allow_missing_keys: bool = False, ) -> None: @@ -829,38 +588,16 @@ def __init__( self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.resizer = Resize(spatial_size=spatial_size, size_mode=size_mode) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, mode, align_corners in self.key_iterator(d, self.mode, self.align_corners): - self.push_transform( - d, - key, - extra_info={ - "mode": mode.value if isinstance(mode, Enum) else mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - ) d[key] = self.resizer(d[key], mode=mode, align_corners=align_corners) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - orig_size = transform[TraceKeys.ORIG_SIZE] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - # Create inverse transform - inverse_transform = Resize( - spatial_size=orig_size, - mode=mode, - align_corners=None if align_corners == TraceKeys.NONE else align_corners, - ) - # Apply inverse transform - d[key] = inverse_transform(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.resizer.inverse(d[key]) return d @@ -881,8 +618,8 @@ def __init__( scale_params: Optional[Union[Sequence[float], float]] = None, affine: Optional[NdarrayOrTensor] = None, spatial_size: Optional[Union[Sequence[int], int]] = None, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.REFLECTION, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.REFLECTION, as_tensor_output: bool = True, device: Optional[torch.device] = None, dtype: Union[DtypeLike, torch.dtype] = np.float32, @@ -952,44 +689,16 @@ def __init__( self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode): - orig_size = d[key].shape[1:] - d[key], affine = self.affine(d[key], mode=mode, padding_mode=padding_mode) - self.push_transform( - d, - key, - orig_size=orig_size, - extra_info={ - "affine": affine, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - }, - ) + d[key], _ = self.affine(d[key], mode=mode, padding_mode=padding_mode) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - orig_size = transform[TraceKeys.ORIG_SIZE] - # Create inverse transform - fwd_affine = transform[TraceKeys.EXTRA_INFO]["affine"] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - inv_affine = np.linalg.inv(fwd_affine) - - affine_grid = AffineGrid(affine=inv_affine) - grid, _ = affine_grid(orig_size) - - # Apply inverse transform - d[key] = self.affine.resampler(d[key], grid, mode, padding_mode) - - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.affine.inverse(d[key]) return d @@ -1010,8 +719,8 @@ def __init__( shear_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, translate_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, scale_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.REFLECTION, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.REFLECTION, cache_grid: bool = False, as_tensor_output: bool = True, device: Optional[torch.device] = None, @@ -1098,64 +807,41 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: - return d + out: Dict[Hashable, NdarrayOrTensor] = convert_to_tensor(d, track_meta=get_track_meta()) + return out self.randomize(None) # all the keys share the same random Affine factor self.rand_affine.randomize() - device = self.rand_affine.resampler.device spatial_size = d[first_key].shape[1:] # type: ignore sp_size = fall_back_tuple(self.rand_affine.spatial_size, spatial_size) # change image size or do random transform do_resampling = self._do_transform or (sp_size != ensure_tuple(spatial_size)) - affine: torch.Tensor = torch.eye(len(sp_size) + 1, dtype=torch.float64, device=device) # converting affine to tensor because the resampler currently only support torch backend grid = None if do_resampling: # need to prepare grid grid = self.rand_affine.get_identity_grid(sp_size) if self._do_transform: # add some random factors grid = self.rand_affine.rand_affine_grid(grid=grid) - affine = self.rand_affine.rand_affine_grid.get_transformation_matrix() for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode): - self.push_transform( - d, - key, - extra_info={ - "affine": affine, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - }, - ) # do the transform if do_resampling: - d[key] = self.rand_affine.resampler(d[key], grid, mode=mode, padding_mode=padding_mode) - + d[key] = self.rand_affine(d[key], mode=mode, padding_mode=padding_mode, grid=grid) + if get_track_meta(): + xform = self.pop_transform(d[key], check=False) if do_resampling else {} + self.push_transform(d[key], extra_info={"do_resampling": do_resampling, "rand_affine_info": xform}) return d def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) - + d = dict(data) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # if transform was not performed and spatial size is None, nothing to do. - if transform[TraceKeys.DO_TRANSFORM] or self.rand_affine.spatial_size is not None: - orig_size = transform[TraceKeys.ORIG_SIZE] - # Create inverse transform - fwd_affine = transform[TraceKeys.EXTRA_INFO]["affine"] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - inv_affine = np.linalg.inv(fwd_affine) - - affine_grid = AffineGrid(affine=inv_affine) - grid, _ = affine_grid(orig_size) - - # Apply inverse transform - d[key] = self.rand_affine.resampler(d[key], grid, mode, padding_mode) - - # Remove the applied transform - self.pop_transform(d, key) + tr = self.pop_transform(d[key]) + do_resampling = tr[TraceKeys.EXTRA_INFO]["do_resampling"] + if do_resampling: + d[key].applied_operations.append(tr[TraceKeys.EXTRA_INFO]["rand_affine_info"]) # type: ignore + d[key] = self.rand_affine.inverse(d[key]) return d @@ -1179,8 +865,8 @@ def __init__( shear_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, translate_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, scale_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.REFLECTION, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.REFLECTION, as_tensor_output: bool = False, device: Optional[torch.device] = None, allow_missing_keys: bool = False, @@ -1266,7 +952,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: - return d + out: Dict[Hashable, NdarrayOrTensor] = convert_to_tensor(d, track_meta=get_track_meta()) + return out self.randomize(None) @@ -1277,7 +964,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N if self._do_transform: grid = self.rand_2d_elastic.deform_grid(spatial_size=sp_size) grid = self.rand_2d_elastic.rand_affine_grid(grid=grid) - grid = torch.nn.functional.interpolate( # type: ignore + grid = torch.nn.functional.interpolate( recompute_scale_factor=True, input=grid.unsqueeze(0), scale_factor=ensure_tuple_rep(self.rand_2d_elastic.deform_grid.spacing, 2), @@ -1313,8 +1000,8 @@ def __init__( shear_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, translate_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, scale_range: Optional[Union[Sequence[Union[Tuple[float, float], float]], float]] = None, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.REFLECTION, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.REFLECTION, as_tensor_output: bool = False, device: Optional[torch.device] = None, allow_missing_keys: bool = False, @@ -1398,11 +1085,12 @@ def set_random_state( super().set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: - return d + out: Dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta()) + return out self.randomize(None) @@ -1448,22 +1136,16 @@ def __init__( super().__init__(keys, allow_missing_keys) self.flipper = Flip(spatial_axis=spatial_axis) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key in self.key_iterator(d): - self.push_transform(d, key) d[key] = self.flipper(d[key]) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - _ = self.get_most_recent_transform(d, key) - # Inverse is same as forward - d[key] = self.flipper(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.flipper.inverse(d[key]) return d @@ -1481,7 +1163,7 @@ class RandFlipd(RandomizableTransform, MapTransform, InvertibleTransform): allow_missing_keys: don't raise exception if key is missing. """ - backend = RandFlip.backend + backend = Flip.backend def __init__( self, @@ -1492,35 +1174,36 @@ def __init__( ) -> None: MapTransform.__init__(self, keys, allow_missing_keys) RandomizableTransform.__init__(self, prob) - self.flipper = RandFlip(prob=1.0, spatial_axis=spatial_axis) + self.flipper = Flip(spatial_axis=spatial_axis) def set_random_state( self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None ) -> "RandFlipd": super().set_random_state(seed, state) - self.flipper.set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) self.randomize(None) for key in self.key_iterator(d): if self._do_transform: - d[key] = self.flipper(d[key], randomize=False) - self.push_transform(d, key) + d[key] = self.flipper(d[key]) + else: + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + if get_track_meta(): + xform_info = self.pop_transform(d[key], check=False) if self._do_transform else {} + self.push_transform(d[key], extra_info=xform_info) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Check if random transform was actually performed (based on `prob`) - if transform[TraceKeys.DO_TRANSFORM]: - # Inverse is same as forward - d[key] = self.flipper(d[key], randomize=False) - # Remove the applied transform - self.pop_transform(d, key) + xform = self.pop_transform(d[key]) + if not xform[TraceKeys.DO_TRANSFORM]: + continue + with self.flipper.trace_transform(False): + d[key] = self.flipper(d[key]) return d @@ -1552,7 +1235,7 @@ def set_random_state( self.flipper.set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: @@ -1565,20 +1248,20 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N for key in self.key_iterator(d): if self._do_transform: d[key] = self.flipper(d[key], randomize=False) - self.push_transform(d, key, extra_info={"axis": self.flipper._axis}) + else: + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + if get_track_meta(): + xform = self.pop_transform(d[key], check=False) if self._do_transform else {} + self.push_transform(d[key], extra_info=xform) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: - d = deepcopy(dict(data)) + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = dict(data) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Check if random transform was actually performed (based on `prob`) - if transform[TraceKeys.DO_TRANSFORM]: - flipper = Flip(spatial_axis=transform[TraceKeys.EXTRA_INFO]["axis"]) - # Inverse is same as forward - d[key] = flipper(d[key]) - # Remove the applied transform - self.pop_transform(d, key) + xform = self.pop_transform(d[key]) + if xform[TraceKeys.DO_TRANSFORM]: + d[key].applied_operations.append(xform[TraceKeys.EXTRA_INFO]) # type: ignore + d[key] = self.flipper.inverse(d[key]) return d @@ -1617,8 +1300,8 @@ def __init__( keys: KeysCollection, angle: Union[Sequence[float], float], keep_size: bool = True, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.BORDER, align_corners: Union[Sequence[bool], bool] = False, dtype: Union[Sequence[Union[DtypeLike, torch.dtype]], DtypeLike, torch.dtype] = np.float32, allow_missing_keys: bool = False, @@ -1631,56 +1314,20 @@ def __init__( self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.dtype = ensure_tuple_rep(dtype, len(self.keys)) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, mode, padding_mode, align_corners, dtype in self.key_iterator( d, self.mode, self.padding_mode, self.align_corners, self.dtype ): - orig_size = d[key].shape[1:] d[key] = self.rotator( d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype ) - rot_mat = self.rotator.get_rotation_matrix() - self.push_transform( - d, - key, - orig_size=orig_size, - extra_info={ - "rot_mat": rot_mat, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - ) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - for key, dtype in self.key_iterator(d, self.dtype): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - fwd_rot_mat = transform[TraceKeys.EXTRA_INFO]["rot_mat"] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - inv_rot_mat = np.linalg.inv(fwd_rot_mat) - - xform = AffineTransform( - normalized=False, - mode=mode, - padding_mode=padding_mode, - align_corners=False if align_corners == TraceKeys.NONE else align_corners, - reverse_indexing=True, - ) - img_t, *_ = convert_data_type(d[key], torch.Tensor, dtype=dtype) - transform_t, *_ = convert_to_dst_type(inv_rot_mat, img_t) - - out = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[TraceKeys.ORIG_SIZE]).squeeze(0) - out, *_ = convert_to_dst_type(out, dst=d[key], dtype=out.dtype) - d[key] = out - # Remove the applied transform - self.pop_transform(d, key) - + for key in self.key_iterator(d): + d[key] = self.rotator.inverse(d[key]) return d @@ -1694,9 +1341,9 @@ class RandRotated(RandomizableTransform, MapTransform, InvertibleTransform): range_x: Range of rotation angle in radians in the plane defined by the first and second axes. If single number, angle is uniformly sampled from (-range_x, range_x). range_y: Range of rotation angle in radians in the plane defined by the first and third axes. - If single number, angle is uniformly sampled from (-range_y, range_y). + If single number, angle is uniformly sampled from (-range_y, range_y). only work for 3D data. range_z: Range of rotation angle in radians in the plane defined by the second and third axes. - If single number, angle is uniformly sampled from (-range_z, range_z). + If single number, angle is uniformly sampled from (-range_z, range_z). only work for 3D data. prob: Probability of rotation. keep_size: If it is False, the output shape is adapted so that the input array is contained completely in the output. @@ -1729,8 +1376,8 @@ def __init__( range_z: Union[Tuple[float, float], float] = 0.0, prob: float = 0.1, keep_size: bool = True, - mode: GridSampleModeSequence = GridSampleMode.BILINEAR, - padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER, + mode: SequenceStr = GridSampleMode.BILINEAR, + padding_mode: SequenceStr = GridSamplePadMode.BORDER, align_corners: Union[Sequence[bool], bool] = False, dtype: Union[Sequence[Union[DtypeLike, torch.dtype]], DtypeLike, torch.dtype] = np.float32, allow_missing_keys: bool = False, @@ -1750,7 +1397,7 @@ def set_random_state( self.rand_rotate.set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) self.randomize(None) @@ -1760,59 +1407,28 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d, self.mode, self.padding_mode, self.align_corners, self.dtype ): if self._do_transform: - d[key], rot_mat = self.rand_rotate( + d[key] = self.rand_rotate( d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners, dtype=dtype, randomize=False, - get_matrix=True, ) else: - rot_mat = np.eye(d[key].ndim) - self.push_transform( - d, - key, - orig_size=d[key].shape[1:], - extra_info={ - "rot_mat": rot_mat, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - ) + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + if get_track_meta(): + rot_info = self.pop_transform(d[key], check=False) if self._do_transform else {} + self.push_transform(d[key], extra_info=rot_info) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) - for key, dtype in self.key_iterator(d, self.dtype): - transform = self.get_most_recent_transform(d, key) - # Check if random transform was actually performed (based on `prob`) - if transform[TraceKeys.DO_TRANSFORM]: - # Create inverse transform - fwd_rot_mat = transform[TraceKeys.EXTRA_INFO]["rot_mat"] - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - inv_rot_mat = np.linalg.inv(fwd_rot_mat) - - xform = AffineTransform( - normalized=False, - mode=mode, - padding_mode=padding_mode, - align_corners=False if align_corners == TraceKeys.NONE else align_corners, - reverse_indexing=True, - ) - img_t, *_ = convert_data_type(d[key], torch.Tensor, dtype=dtype) - transform_t, *_ = convert_to_dst_type(inv_rot_mat, img_t) - output: torch.Tensor - out = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[TraceKeys.ORIG_SIZE]).squeeze(0) - out, *_ = convert_to_dst_type(out, dst=d[key], dtype=out.dtype) - d[key] = out - # Remove the applied transform - self.pop_transform(d, key) - + for key in self.key_iterator(d): + xform = self.pop_transform(d[key]) + if xform[TraceKeys.DO_TRANSFORM]: + d[key].applied_operations.append(xform[TraceKeys.EXTRA_INFO]) # type: ignore + d[key] = self.rand_rotate.inverse(d[key]) return d @@ -1825,7 +1441,7 @@ class Zoomd(MapTransform, InvertibleTransform): zoom: The zoom factor along the spatial axes. If a float, zoom is the same for each spatial axis. If a sequence, zoom should contain one value for each spatial axis. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html It also can be a sequence of string, each element corresponds to a key in ``keys``. @@ -1853,8 +1469,8 @@ def __init__( self, keys: KeysCollection, zoom: Union[Sequence[float], float], - mode: InterpolateModeSequence = InterpolateMode.AREA, - padding_mode: PadModeSequence = NumpyPadMode.EDGE, + mode: SequenceStr = InterpolateMode.AREA, + padding_mode: SequenceStr = NumpyPadMode.EDGE, align_corners: Union[Sequence[Optional[bool]], Optional[bool]] = None, keep_size: bool = True, allow_missing_keys: bool = False, @@ -1866,45 +1482,18 @@ def __init__( self.align_corners = ensure_tuple_rep(align_corners, len(self.keys)) self.zoomer = Zoom(zoom=zoom, keep_size=keep_size, **kwargs) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, mode, padding_mode, align_corners in self.key_iterator( d, self.mode, self.padding_mode, self.align_corners ): - self.push_transform( - d, - key, - extra_info={ - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - ) d[key] = self.zoomer(d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Create inverse transform - zoom = np.array(self.zoomer.zoom) - inverse_transform = Zoom(zoom=(1 / zoom).tolist(), keep_size=self.zoomer.keep_size) - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - # Apply inverse - d[key] = inverse_transform( - d[key], - mode=mode, - padding_mode=padding_mode, - align_corners=None if align_corners == TraceKeys.NONE else align_corners, - ) - # Size might be out by 1 voxel so pad - d[key] = SpatialPad(transform[TraceKeys.ORIG_SIZE], mode="edge")(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + d[key] = self.zoomer.inverse(d[key]) return d @@ -1925,7 +1514,7 @@ class RandZoomd(RandomizableTransform, MapTransform, InvertibleTransform): to keep the original spatial shape ratio. If a sequence, max_zoom should contain one value for each spatial axis. If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio. - mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} + mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``} The interpolation mode. Defaults to ``"area"``. See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html It also can be a sequence of string, each element corresponds to a key in ``keys``. @@ -1955,8 +1544,8 @@ def __init__( prob: float = 0.1, min_zoom: Union[Sequence[float], float] = 0.9, max_zoom: Union[Sequence[float], float] = 1.1, - mode: InterpolateModeSequence = InterpolateMode.AREA, - padding_mode: PadModeSequence = NumpyPadMode.EDGE, + mode: SequenceStr = InterpolateMode.AREA, + padding_mode: SequenceStr = NumpyPadMode.EDGE, align_corners: Union[Sequence[Optional[bool]], Optional[bool]] = None, keep_size: bool = True, allow_missing_keys: bool = False, @@ -1976,11 +1565,12 @@ def set_random_state( self.rand_zoom.set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: - return d + out: Dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta()) + return out self.randomize(None) @@ -1993,42 +1583,20 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d[key] = self.rand_zoom( d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners, randomize=False ) - self.push_transform( - d, - key, - extra_info={ - "zoom": self.rand_zoom._zoom, - "mode": mode.value if isinstance(mode, Enum) else mode, - "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode, - "align_corners": align_corners if align_corners is not None else TraceKeys.NONE, - }, - ) + else: + d[key] = convert_to_tensor(d[key], track_meta=get_track_meta()) + if get_track_meta(): + xform = self.pop_transform(d[key], check=False) if self._do_transform else {} + self.push_transform(d[key], extra_info=xform) return d - def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = deepcopy(dict(data)) for key in self.key_iterator(d): - transform = self.get_most_recent_transform(d, key) - # Check if random transform was actually performed (based on `prob`) - if transform[TraceKeys.DO_TRANSFORM]: - # Create inverse transform - zoom = np.array(transform[TraceKeys.EXTRA_INFO]["zoom"]) - mode = transform[TraceKeys.EXTRA_INFO]["mode"] - padding_mode = transform[TraceKeys.EXTRA_INFO]["padding_mode"] - align_corners = transform[TraceKeys.EXTRA_INFO]["align_corners"] - inverse_transform = Zoom(zoom=(1 / zoom).tolist(), keep_size=self.rand_zoom.keep_size) - # Apply inverse - d[key] = inverse_transform( - d[key], - mode=mode, - padding_mode=padding_mode, - align_corners=None if align_corners == TraceKeys.NONE else align_corners, - ) - # Size might be out by 1 voxel so pad - d[key] = SpatialPad(transform[TraceKeys.ORIG_SIZE], mode="edge")(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - + xform = self.pop_transform(d[key]) + if xform[TraceKeys.DO_TRANSFORM]: + d[key].applied_operations.append(xform[TraceKeys.EXTRA_INFO]) # type: ignore + d[key] = self.rand_zoom.inverse(d[key]) return d @@ -2044,8 +1612,8 @@ def __init__( keys: KeysCollection, num_cells: Union[Tuple[int], int], distort_steps: List[Tuple], - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, device: Optional[torch.device] = None, allow_missing_keys: bool = False, ) -> None: @@ -2073,7 +1641,7 @@ def __init__( self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode): d[key] = self.grid_distortion(d[key], mode=mode, padding_mode=padding_mode) @@ -2093,8 +1661,8 @@ def __init__( num_cells: Union[Tuple[int], int] = 5, prob: float = 0.1, distort_limit: Union[Tuple[float, float], float] = (-0.03, 0.03), - mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR, - padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER, + mode: str = GridSampleMode.BILINEAR, + padding_mode: str = GridSamplePadMode.BORDER, device: Optional[torch.device] = None, allow_missing_keys: bool = False, ) -> None: @@ -2133,15 +1701,17 @@ def set_random_state( self.rand_grid_distortion.set_random_state(seed, state) return self - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) self.randomize(None) if not self._do_transform: - return d + out: Dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta()) + return out first_key: Union[Hashable, List] = self.first_key(d) if first_key == []: - return d + out = convert_to_tensor(d, track_meta=get_track_meta()) + return out self.rand_grid_distortion.randomize(d[first_key].shape[1:]) # type: ignore for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode): @@ -2149,6 +1719,222 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d +class GridSplitd(MapTransform): + """ + Split the image into patches based on the provided grid in 2D. + + Args: + keys: keys of the corresponding items to be transformed. + grid: a tuple define the shape of the grid upon which the image is split. Defaults to (2, 2) + size: a tuple or an integer that defines the output patch sizes, + or a dictionary that define it separately for each key, like {"image": 3, "mask", (2, 2)}. + If it's an integer, the value will be repeated for each dimension. + The default is None, where the patch size will be inferred from the grid shape. + allow_missing_keys: don't raise exception if key is missing. + + Note: This transform currently support only image with two spatial dimensions. + """ + + backend = GridSplit.backend + + def __init__( + self, + keys: KeysCollection, + grid: Tuple[int, int] = (2, 2), + size: Optional[Union[int, Tuple[int, int], Dict[Hashable, Union[int, Tuple[int, int], None]]]] = None, + allow_missing_keys: bool = False, + ): + super().__init__(keys, allow_missing_keys) + self.grid = grid + self.size = size if isinstance(size, dict) else {key: size for key in self.keys} + self.splitter = GridSplit(grid=grid) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict[Hashable, NdarrayOrTensor]]: + d = dict(data) + n_outputs = np.prod(self.grid) + output: List[Dict[Hashable, NdarrayOrTensor]] = [dict(d) for _ in range(n_outputs)] + for key in self.key_iterator(d): + result = self.splitter(d[key], self.size[key]) + for i in range(n_outputs): + output[i][key] = result[i] + return output + + +class GridPatchd(MapTransform): + """ + Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps. + It can sort the patches and return all or a subset of them. + + Args: + keys: keys of the corresponding items to be transformed. + patch_size: size of patches to generate slices for, 0 or None selects whole dimension + offset: starting position in the array, default is 0 for each dimension. + np.random.randint(0, patch_size, 2) creates random start between 0 and `patch_size` for a 2D image. + num_patches: number of patches to return. Defaults to None, which returns all the available patches. + overlap: amount of overlap between patches in each dimension. Default to 0.0. + sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`), + lowest values (`"min"`), or in their default order (`None`). Default to None. + threshold: a value to keep only the patches whose sum of intensities are less than the threshold. + Defaults to no filtering. + pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``. + allow_missing_keys: don't raise exception if key is missing. + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + + Returns: + a list of dictionaries, each of which contains the all the original key/value with the values for `keys` + replaced by the patches. It also add the following new keys: + + "patch_location": the starting location of the patch in the image, + "patch_size": size of the extracted patch + "num_patches": total number of patches in the image + "offset": the amount of offset for the patches in the image (starting position of upper left patch) + """ + + backend = GridPatch.backend + + def __init__( + self, + keys: KeysCollection, + patch_size: Sequence[int], + offset: Optional[Sequence[int]] = None, + num_patches: Optional[int] = None, + overlap: float = 0.0, + sort_fn: Optional[str] = None, + threshold: Optional[float] = None, + pad_mode: str = PytorchPadMode.CONSTANT, + allow_missing_keys: bool = False, + **pad_kwargs, + ): + super().__init__(keys, allow_missing_keys) + self.patcher = GridPatch( + patch_size=patch_size, + offset=offset, + num_patches=num_patches, + overlap=overlap, + sort_fn=sort_fn, + threshold=threshold, + pad_mode=pad_mode, + **pad_kwargs, + ) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict]: + d = dict(data) + original_spatial_shape = d[first(self.keys)].shape[1:] + output = [] + results = [self.patcher(d[key]) for key in self.keys] + num_patches = min(len(r) for r in results) + for patch in zip(*results): + new_dict = {k: v[0] for k, v in zip(self.keys, patch)} + # fill in the extra keys with unmodified data + for k in set(d.keys()).difference(set(self.keys)): + new_dict[k] = deepcopy(d[k]) + # fill additional metadata + new_dict["original_spatial_shape"] = original_spatial_shape + new_dict[WSIPatchKeys.LOCATION] = patch[0][1] # use the starting coordinate of the first item + new_dict[WSIPatchKeys.SIZE] = self.patcher.patch_size + new_dict[WSIPatchKeys.COUNT] = num_patches + new_dict["offset"] = self.patcher.offset + output.append(new_dict) + return output + + +class RandGridPatchd(RandomizableTransform, MapTransform): + """ + Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps, + and with random offset for the minimal corner of the image, (0,0) for 2D and (0,0,0) for 3D. + It can sort the patches and return all or a subset of them. + + Args: + keys: keys of the corresponding items to be transformed. + patch_size: size of patches to generate slices for, 0 or None selects whole dimension + min_offset: the minimum range of starting position to be selected randomly. Defaults to 0. + max_offset: the maximum range of starting position to be selected randomly. + Defaults to image size modulo patch size. + num_patches: number of patches to return. Defaults to None, which returns all the available patches. + overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0). + If only one float number is given, it will be applied to all dimensions. Defaults to 0.0. + sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`), + lowest values (`"min"`), or in their default order (`None`). Default to None. + threshold: a value to keep only the patches whose sum of intensities are less than the threshold. + Defaults to no filtering. + pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``. + allow_missing_keys: don't raise exception if key is missing. + pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. + + Returns: + a list of dictionaries, each of which contains the all the original key/value with the values for `keys` + replaced by the patches. It also add the following new keys: + + "patch_location": the starting location of the patch in the image, + "patch_size": size of the extracted patch + "num_patches": total number of patches in the image + "offset": the amount of offset for the patches in the image (starting position of the first patch) + + """ + + backend = RandGridPatch.backend + + def __init__( + self, + keys: KeysCollection, + patch_size: Sequence[int], + min_offset: Optional[Union[Sequence[int], int]] = None, + max_offset: Optional[Union[Sequence[int], int]] = None, + num_patches: Optional[int] = None, + overlap: float = 0.0, + sort_fn: Optional[str] = None, + threshold: Optional[float] = None, + pad_mode: str = PytorchPadMode.CONSTANT, + allow_missing_keys: bool = False, + **pad_kwargs, + ): + MapTransform.__init__(self, keys, allow_missing_keys) + self.patcher = RandGridPatch( + patch_size=patch_size, + min_offset=min_offset, + max_offset=max_offset, + num_patches=num_patches, + overlap=overlap, + sort_fn=sort_fn, + threshold=threshold, + pad_mode=pad_mode, + **pad_kwargs, + ) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandGridPatchd": + super().set_random_state(seed, state) + self.patcher.set_random_state(seed, state) + return self + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> List[Dict]: + d = dict(data) + original_spatial_shape = d[first(self.keys)].shape[1:] + # all the keys share the same random noise + first_key: Union[Hashable, List] = self.first_key(d) + if first_key == []: + return [d] + self.patcher.randomize(d[first_key]) # type: ignore + results = [self.patcher(d[key], randomize=False) for key in self.keys] + + num_patches = min(len(r) for r in results) + output = [] + for patch in zip(*results): + new_dict = {k: v[0] for k, v in zip(self.keys, patch)} + # fill in the extra keys with unmodified data + for k in set(d.keys()).difference(set(self.keys)): + new_dict[k] = deepcopy(d[k]) + # fill additional metadata + new_dict["original_spatial_shape"] = original_spatial_shape + new_dict[WSIPatchKeys.LOCATION] = patch[0][1] # use the starting coordinate of the first item + new_dict[WSIPatchKeys.SIZE] = self.patcher.patch_size + new_dict[WSIPatchKeys.COUNT] = num_patches + new_dict["offset"] = self.patcher.offset + output.append(new_dict) + return output + + SpatialResampleD = SpatialResampleDict = SpatialResampled ResampleToMatchD = ResampleToMatchDict = ResampleToMatchd SpacingD = SpacingDict = Spacingd @@ -2169,3 +1955,6 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N RandRotateD = RandRotateDict = RandRotated ZoomD = ZoomDict = Zoomd RandZoomD = RandZoomDict = RandZoomd +GridSplitD = GridSplitDict = GridSplitd +GridPatchD = GridPatchDict = GridPatchd +RandGridPatchD = RandGridPatchDict = RandGridPatchd diff --git a/monai/transforms/transform.py b/monai/transforms/transform.py index 8537f7eb89..5819d2971d 100644 --- a/monai/transforms/transform.py +++ b/monai/transforms/transform.py @@ -14,7 +14,7 @@ import logging from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, Generator, Hashable, Iterable, List, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Dict, Generator, Hashable, Iterable, List, Mapping, Optional, Tuple, TypeVar, Union import numpy as np import torch @@ -75,7 +75,7 @@ def apply_transform( unpack_items: whether to unpack parameters using `*`. Defaults to False. log_stats: whether to log the detailed information of data and applied transform when error happened, for NumPy array and PyTorch Tensor, log the data shape and value range, - for other meta data, log the values directly. default to `False`. + for other metadata, log the values directly. default to `False`. Raises: Exception: When ``transform`` raises an exception. @@ -102,7 +102,7 @@ def _log_stats(data, prefix: Optional[str] = "Data"): # log data type, shape, range for array datastats(img=data, data_shape=True, value_range=True, prefix=prefix) else: - # log data type and value for other meta data + # log data type and value for other metadata datastats(img=data, data_value=True, prefix=prefix) if isinstance(data, dict): @@ -348,7 +348,7 @@ def __call__(self, data): """ raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.") - def key_iterator(self, data: Dict[Hashable, Any], *extra_iterables: Optional[Iterable]) -> Generator: + def key_iterator(self, data: Mapping[Hashable, Any], *extra_iterables: Optional[Iterable]) -> Generator: """ Iterate across keys and optionally extra iterables. If key is missing, exception is raised if `allow_missing_keys==False` (default). If `allow_missing_keys==True`, key is skipped. diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 72f50f6894..41973ccb47 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -17,6 +17,7 @@ import sys import time import warnings +from copy import deepcopy from typing import Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union import numpy as np @@ -24,6 +25,10 @@ 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.utils import no_collation +from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import Randomizable, RandomizableTransform, Transform from monai.transforms.utils import ( extreme_points_to_image, @@ -33,10 +38,12 @@ ) from monai.transforms.utils_pytorch_numpy_unification import concatenate, in1d, moveaxis, unravel_indices from monai.utils import ( + TraceKeys, convert_data_type, convert_to_cupy, convert_to_numpy, convert_to_tensor, + deprecated, deprecated_arg, ensure_tuple, look_up_option, @@ -62,6 +69,7 @@ "EnsureType", "RepeatChannel", "RemoveRepeatedChannel", + "SplitDim", "SplitChannel", "CastToType", "ToTensor", @@ -131,7 +139,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ - return moveaxis(img, self.channel_dim, 0) + out: NdarrayOrTensor = convert_to_tensor(moveaxis(img, self.channel_dim, 0), track_meta=get_track_meta()) + return out class AsChannelLast(Transform): @@ -160,7 +169,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ - return moveaxis(img, self.channel_dim, -1) + out: NdarrayOrTensor = convert_to_tensor(moveaxis(img, self.channel_dim, -1), track_meta=get_track_meta()) + return out class AddChannel(Transform): @@ -183,7 +193,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ - return img[None] + out: NdarrayOrTensor = convert_to_tensor(img[None], track_meta=get_track_meta()) + return out class EnsureChannelFirst(Transform): @@ -204,18 +215,19 @@ def __init__(self, strict_check: bool = True): self.strict_check = strict_check self.add_channel = AddChannel() - def __call__(self, img: NdarrayOrTensor, meta_dict: Optional[Mapping] = None) -> NdarrayOrTensor: + def __call__(self, img: torch.Tensor, meta_dict: Optional[Mapping] = None) -> torch.Tensor: """ Apply the transform to `img`. """ - if not isinstance(meta_dict, Mapping): - msg = "meta_dict not available, EnsureChannelFirst is not in use." + if not isinstance(img, MetaTensor) and not isinstance(meta_dict, Mapping): + msg = "metadata not available, EnsureChannelFirst is not in use." if self.strict_check: raise ValueError(msg) warnings.warn(msg) return img - - channel_dim = meta_dict.get("original_channel_dim") + if isinstance(img, MetaTensor): + meta_dict = img.meta + channel_dim = meta_dict.get("original_channel_dim") # type: ignore if channel_dim is None: msg = "Unknown original_channel_dim in the meta_dict, EnsureChannelFirst is not in use." @@ -224,8 +236,8 @@ def __call__(self, img: NdarrayOrTensor, meta_dict: Optional[Mapping] = None) -> warnings.warn(msg) return img if channel_dim == "no_channel": - return self.add_channel(img) - return AsChannelFirst(channel_dim=channel_dim)(img) + return self.add_channel(img) # type: ignore + return AsChannelFirst(channel_dim=channel_dim)(img) # type: ignore class RepeatChannel(Transform): @@ -250,7 +262,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Apply the transform to `img`, assuming `img` is a "channel-first" array. """ repeat_fn = torch.repeat_interleave if isinstance(img, torch.Tensor) else np.repeat - return repeat_fn(img, self.repeats, 0) # type: ignore + return convert_to_tensor(repeat_fn(img, self.repeats, 0), track_meta=get_track_meta()) # type: ignore class RemoveRepeatedChannel(Transform): @@ -278,36 +290,73 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: if img.shape[0] < 2: raise AssertionError("Image must have more than one channel") - return img[:: self.repeats, :] + out: NdarrayOrTensor = convert_to_tensor(img[:: self.repeats, :], track_meta=get_track_meta()) + return out -class SplitChannel(Transform): +class SplitDim(Transform): """ - Split Numpy array or PyTorch Tensor data according to the channel dim. - It can help applying different following transforms to different channels. + Given an image of size X along a certain dimension, return a list of length X containing + images. Useful for converting 3D images into a stack of 2D images, splitting multichannel inputs into + single channels, for example. - Args: - channel_dim: which dimension of input image is the channel, default to 0. + Note: `torch.split`/`np.split` is used, so the outputs are views of the input (shallow copy). + Args: + dim: dimension on which to split + keepdim: if `True`, output will have singleton in the split dimension. If `False`, this + dimension will be squeezed. + update_meta: whether to update the MetaObj in each split result. """ backend = [TransformBackends.TORCH, TransformBackends.NUMPY] - def __init__(self, channel_dim: int = 0) -> None: - self.channel_dim = channel_dim + def __init__(self, dim: int = -1, keepdim: bool = True, update_meta=True) -> None: + self.dim = dim + self.keepdim = keepdim + self.update_meta = update_meta - def __call__(self, img: NdarrayOrTensor) -> List[NdarrayOrTensor]: - num_classes = img.shape[self.channel_dim] - if num_classes <= 1: - raise RuntimeError("input image does not contain multiple channels.") + def __call__(self, img: torch.Tensor) -> List[torch.Tensor]: + """ + Apply the transform to `img`. + """ + n_out = img.shape[self.dim] + if n_out <= 1: + raise RuntimeError(f"Input image is singleton along dimension to be split, got shape {img.shape}.") + if isinstance(img, torch.Tensor): + outputs = list(torch.split(img, 1, self.dim)) + else: + outputs = np.split(img, n_out, self.dim) + for idx, item in enumerate(outputs): + if not self.keepdim: + outputs[idx] = item.squeeze(self.dim) + if self.update_meta and isinstance(img, MetaTensor): + if not isinstance(item, MetaTensor): + item = MetaTensor(item, meta=deepcopy(img.meta)) + if self.dim == 0: # don't update affine if channel dim + 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 + return outputs - outputs = [] - slices = [slice(None)] * len(img.shape) - for i in range(num_classes): - slices[self.channel_dim] = slice(i, i + 1) - outputs.append(img[tuple(slices)]) - return outputs +@deprecated(since="0.8", msg_suffix="please use `SplitDim` instead.") +class SplitChannel(SplitDim): + """ + Split Numpy array or PyTorch Tensor data according to the channel dim. + It can help applying different following transforms to different channels. + + Note: `torch.split`/`np.split` is used, so the outputs are views of the input (shallow copy). + + Args: + channel_dim: which dimension of input image is the channel, default to 0. + + """ + + def __init__(self, channel_dim: int = 0) -> None: + super().__init__(channel_dim) class CastToType(Transform): @@ -368,6 +417,8 @@ def __call__(self, img: NdarrayOrTensor): """ Apply the transform to `img` and make it contiguous. """ + if isinstance(img, MetaTensor): + img.applied_operations = [] # drops tracking info return convert_to_tensor(img, dtype=self.dtype, device=self.device, wrap_sequence=self.wrap_sequence) @@ -384,6 +435,8 @@ class EnsureType(Transform): device: for Tensor data type, specify the target device. wrap_sequence: if `False`, then lists will recursively call this function, default to `True`. E.g., if `False`, `[1, 2]` -> `[tensor(1), tensor(2)]`, if `True`, then `[1, 2]` -> `tensor([1, 2])`. + track_meta: whether to convert to `MetaTensor` when `data_type` is "tensor". + If False, the output data type will be `torch.Tensor`. Default to the return value of ``get_track_meta``. """ @@ -395,11 +448,13 @@ def __init__( dtype: Optional[Union[DtypeLike, torch.dtype]] = None, device: Optional[torch.device] = None, wrap_sequence: bool = True, + track_meta: Optional[bool] = None, ) -> None: self.data_type = look_up_option(data_type.lower(), {"tensor", "numpy"}) self.dtype = dtype self.device = device self.wrap_sequence = wrap_sequence + self.track_meta = get_track_meta() if track_meta is None else bool(track_meta) def __call__(self, data: NdarrayOrTensor): """ @@ -410,10 +465,17 @@ def __call__(self, data: NdarrayOrTensor): if applicable and `wrap_sequence=False`. """ - output_type = torch.Tensor if self.data_type == "tensor" else np.ndarray + if self.data_type == "tensor": + output_type = MetaTensor if self.track_meta else torch.Tensor + else: + output_type = np.ndarray # type: ignore out: NdarrayOrTensor out, *_ = convert_data_type( - data=data, output_type=output_type, dtype=self.dtype, device=self.device, wrap_sequence=self.wrap_sequence + data=data, + output_type=output_type, # type: ignore + dtype=self.dtype, + device=self.device, + wrap_sequence=self.wrap_sequence, ) return out @@ -502,9 +564,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ - if isinstance(img, torch.Tensor): - return img.permute(self.indices or tuple(range(img.ndim)[::-1])) - return img.transpose(self.indices) # type: ignore + img = convert_to_tensor(img, track_meta=get_track_meta()) + return img.permute(self.indices or tuple(range(img.ndim)[::-1])) # type: ignore class SqueezeDim(Transform): @@ -533,11 +594,12 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Args: img: numpy arrays with required dimension `dim` removed """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if self.dim is None: return img.squeeze() # for pytorch/numpy unification if img.shape[self.dim] != 1: - raise ValueError("Can only squeeze singleton dimension") + raise ValueError(f"Can only squeeze singleton dimension, got shape {img.shape}.") return img.squeeze(self.dim) @@ -595,10 +657,17 @@ def __init__( _logger = logging.getLogger(self._logger_name) _logger.setLevel(logging.INFO) if logging.root.getEffectiveLevel() > logging.INFO: - # if the root log level is higher than INFO, set a separate stream handler to record - console = logging.StreamHandler(sys.stdout) - console.setLevel(logging.INFO) - _logger.addHandler(console) + # Avoid duplicate stream handlers to be added when multiple DataStats are used in a chain. + has_console_handler = any( + hasattr(h, "is_data_stats_handler") and h.is_data_stats_handler # type:ignore[attr-defined] + for h in _logger.handlers + ) + if not has_console_handler: + # if the root log level is higher than INFO, set a separate stream handler to record + console = logging.StreamHandler(sys.stdout) + console.setLevel(logging.INFO) + console.is_data_stats_handler = True # type:ignore[attr-defined] + _logger.addHandler(console) def __call__( self, @@ -671,7 +740,7 @@ def __call__(self, img: NdarrayOrTensor, delay_time: Optional[float] = None) -> return img -class Lambda(Transform): +class Lambda(InvertibleTransform): """ Apply a user-defined lambda as a transform. @@ -687,6 +756,7 @@ class Lambda(Transform): Args: func: Lambda/function to be applied. + inv_func: Lambda/function of inverse operation, default to `lambda x: x`. Raises: TypeError: When ``func`` is not an ``Optional[Callable]``. @@ -695,10 +765,11 @@ class Lambda(Transform): backend = [TransformBackends.TORCH, TransformBackends.NUMPY] - def __init__(self, func: Optional[Callable] = None) -> None: + def __init__(self, func: Optional[Callable] = None, inv_func: Callable = no_collation) -> None: if func is not None and not callable(func): raise TypeError(f"func must be None or callable but is {type(func).__name__}.") self.func = func + self.inv_func = inv_func def __call__(self, img: NdarrayOrTensor, func: Optional[Callable] = None): """ @@ -709,16 +780,23 @@ def __call__(self, img: NdarrayOrTensor, func: Optional[Callable] = None): Raises: TypeError: When ``func`` is not an ``Optional[Callable]``. - ValueError: When ``func=None`` and ``self.func=None``. Incompatible values. """ - if func is not None: - if not callable(func): - raise TypeError(f"func must be None or callable but is {type(func).__name__}.") - return func(img) - if self.func is not None: - return self.func(img) - raise ValueError("Incompatible values: func=None and self.func=None.") + fn = func if func is not None else self.func + if not callable(fn): + raise TypeError(f"func must be None or callable but is {type(fn).__name__}.") + out = fn(img) + # convert to MetaTensor if necessary + if isinstance(out, (np.ndarray, torch.Tensor)) and not isinstance(out, MetaTensor) and get_track_meta(): + out = MetaTensor(out) + if isinstance(out, MetaTensor): + self.push_transform(out) + return out + + def inverse(self, data: torch.Tensor): + if isinstance(data, MetaTensor): + self.pop_transform(data) + return self.inv_func(data) class RandLambda(Lambda, RandomizableTransform): @@ -729,19 +807,35 @@ class RandLambda(Lambda, RandomizableTransform): Args: func: Lambda/function to be applied. prob: probability of executing the random function, default to 1.0, with 100% probability to execute. + inv_func: Lambda/function of inverse operation, default to `lambda x: x`. For more details, please check :py:class:`monai.transforms.Lambda`. """ backend = Lambda.backend - def __init__(self, func: Optional[Callable] = None, prob: float = 1.0) -> None: - Lambda.__init__(self=self, func=func) + def __init__(self, func: Optional[Callable] = None, prob: float = 1.0, inv_func: Callable = no_collation) -> None: + Lambda.__init__(self=self, func=func, inv_func=inv_func) RandomizableTransform.__init__(self=self, prob=prob) def __call__(self, img: NdarrayOrTensor, func: Optional[Callable] = None): self.randomize(img) - return super().__call__(img=img, func=func) if self._do_transform else img + out = deepcopy(super().__call__(img, func) if self._do_transform else img) + # convert to MetaTensor if necessary + if not isinstance(out, MetaTensor) and get_track_meta(): + out = MetaTensor(out) + if isinstance(out, MetaTensor): + lambda_info = self.pop_transform(out) if self._do_transform else {} + self.push_transform(out, extra_info=lambda_info) + return out + + def inverse(self, data: torch.Tensor): + do_transform = self.get_most_recent_transform(data).pop(TraceKeys.DO_TRANSFORM) + if do_transform: + data = super().inverse(data) + else: + self.pop_transform(data) + return data class LabelToMask(Transform): @@ -785,6 +879,7 @@ def __call__( merge_channels: whether to use `np.any()` to merge the result on channel dim. if yes, will return a single channel mask with binary data. """ + img = convert_to_tensor(img, track_meta=get_track_meta()) if select_labels is None: select_labels = self.select_labels else: @@ -1025,6 +1120,7 @@ def __call__(self, img: NdarrayOrTensor): """ img_t, *_ = convert_data_type(img, torch.Tensor) + out = self.trans(img_t) out, *_ = convert_to_dst_type(src=out, dst=img) return out @@ -1079,7 +1175,7 @@ def __call__(self, img: NdarrayOrTensor): class IntensityStats(Transform): """ - Compute statistics for the intensity values of input image and store into the meta data dictionary. + Compute statistics for the intensity values of input image and store into the metadata dictionary. For example: if `ops=[lambda x: np.mean(x), "max"]` and `key_prefix="orig"`, may generate below stats: `{"orig_custom_0": 1.5, "orig_max": 3.0}`. @@ -1089,7 +1185,7 @@ class IntensityStats(Transform): mapping to `np.nanmean`, `np.nanmedian`, `np.nanmax`, `np.nanmin`, `np.nanstd`. if a callable function, will execute the function on input image. key_prefix: the prefix to combine with `ops` name to generate the key to store the results in the - meta data dictionary. if some `ops` are callable functions, will use "{key_prefix}_custom_{index}" + metadata dictionary. if some `ops` are callable functions, will use "{key_prefix}_custom_{index}" as the key, where index counts from 0. channel_wise: whether to compute statistics for every channel of input image separately. if True, return a list of values for every operation, default to False. @@ -1111,7 +1207,7 @@ def __call__( Args: img: input image to compute intensity stats. - meta_data: meta data dictionary to store the statistics data, if None, will create an empty dictionary. + meta_data: metadata dictionary to store the statistics data, if None, will create an empty dictionary. mask: if not None, mask the image to extract only the interested area to compute statistics. mask must have the same shape as input `img`. diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index ecf9aaffa4..f8b5bb737b 100644 --- a/monai/transforms/utility/dictionary.py +++ b/monai/transforms/utility/dictionary.py @@ -24,6 +24,7 @@ from monai.config import DtypeLike, KeysCollection from monai.config.type_definitions import NdarrayOrTensor +from monai.data.meta_tensor import MetaTensor from monai.data.utils import no_collation from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import MapTransform, Randomizable, RandomizableTransform @@ -49,7 +50,7 @@ RemoveRepeatedChannel, RepeatChannel, SimulateDelay, - SplitChannel, + SplitDim, SqueezeDim, ToCupy, ToDevice, @@ -61,7 +62,7 @@ ) from monai.transforms.utils import extreme_points_to_image, get_extreme_points from monai.transforms.utils_pytorch_numpy_unification import concatenate -from monai.utils import convert_to_numpy, deprecated_arg, ensure_tuple, ensure_tuple_rep +from monai.utils import deprecated, deprecated_arg, ensure_tuple, ensure_tuple_rep from monai.utils.enums import PostFix, TraceKeys, TransformBackends from monai.utils.type_conversion import convert_to_dst_type @@ -150,6 +151,9 @@ "SplitChannelD", "SplitChannelDict", "SplitChanneld", + "SplitDimD", + "SplitDimDict", + "SplitDimd", "SqueezeDimD", "SqueezeDimDict", "SqueezeDimd", @@ -288,6 +292,8 @@ class EnsureChannelFirstd(MapTransform): backend = EnsureChannelFirst.backend + @deprecated_arg(name="meta_keys", since="0.9", msg_suffix="not needed if image is type `MetaTensor`.") + @deprecated_arg(name="meta_key_postfix", since="0.9", msg_suffix="not needed if image is type `MetaTensor`.") def __init__( self, keys: KeysCollection, @@ -299,14 +305,6 @@ def __init__( Args: keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. - for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. - it can be a sequence of string, map to the `keys`. - if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None and `key_{postfix}` was used to store the metadata in `LoadImaged`. - So need the key to extract metadata for channel dim information, default is `meta_dict`. - For example, for data with key `image`, metadata by default is in `image_meta_dict`. strict_check: whether to raise an error when the meta information is insufficient. """ @@ -315,10 +313,10 @@ def __init__( self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys)) self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) - def __call__(self, data) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, meta_key, meta_key_postfix in zip(self.keys, self.meta_keys, self.meta_key_postfix): - d[key] = self.adjuster(d[key], d[meta_key or f"{key}_{meta_key_postfix}"]) + d[key] = self.adjuster(d[key], d.get(meta_key or f"{key}_{meta_key_postfix}")) # type: ignore return d @@ -372,19 +370,14 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d -class SplitChanneld(MapTransform): - """ - Dictionary-based wrapper of :py:class:`monai.transforms.SplitChannel`. - All the input specified by `keys` should be split into same count of data. - """ - - backend = SplitChannel.backend - +class SplitDimd(MapTransform): def __init__( self, keys: KeysCollection, output_postfixes: Optional[Sequence[str]] = None, - channel_dim: int = 0, + dim: int = 0, + keepdim: bool = True, + update_meta: bool = True, allow_missing_keys: bool = False, ) -> None: """ @@ -395,21 +388,24 @@ def __init__( for example: if the key of input data is `pred` and split 2 classes, the output data keys will be: pred_(output_postfixes[0]), pred_(output_postfixes[1]) if None, using the index number: `pred_0`, `pred_1`, ... `pred_N`. - channel_dim: which dimension of input image is the channel, default to 0. + dim: which dimension of input image is the channel, default to 0. + keepdim: if `True`, output will have singleton in the split dimension. If `False`, this + dimension will be squeezed. + update_meta: if `True`, copy `[key]_meta_dict` for each output and update affine to + reflect the cropped image allow_missing_keys: don't raise exception if key is missing. - """ super().__init__(keys, allow_missing_keys) self.output_postfixes = output_postfixes - self.splitter = SplitChannel(channel_dim=channel_dim) + self.splitter = SplitDim(dim, keepdim, update_meta) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key in self.key_iterator(d): rets = self.splitter(d[key]) postfixes: Sequence = list(range(len(rets))) if self.output_postfixes is None else self.output_postfixes if len(postfixes) != len(rets): - raise AssertionError("count of split results must match output_postfixes.") + raise ValueError(f"count of splits must match output_postfixes, {len(postfixes)} != {len(rets)}.") for i, r in enumerate(rets): split_key = f"{key}_{postfixes[i]}" if split_key in d: @@ -418,6 +414,29 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N return d +@deprecated(since="0.8", msg_suffix="please use `SplitDimd` instead.") +class SplitChanneld(SplitDimd): + """ + Dictionary-based wrapper of :py:class:`monai.transforms.SplitChannel`. + All the input specified by `keys` should be split into same count of data. + """ + + def __init__( + self, + keys: KeysCollection, + output_postfixes: Optional[Sequence[str]] = None, + channel_dim: int = 0, + allow_missing_keys: bool = False, + ) -> None: + super().__init__( + keys, + output_postfixes=output_postfixes, + dim=channel_dim, + update_meta=False, # for backwards compatibility + allow_missing_keys=allow_missing_keys, + ) + + class CastToTyped(MapTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.CastToType`. @@ -500,7 +519,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd return d -class EnsureTyped(MapTransform, InvertibleTransform): +class EnsureTyped(MapTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.EnsureType`. @@ -522,6 +541,7 @@ def __init__( dtype: Union[DtypeLike, torch.dtype] = None, device: Optional[torch.device] = None, wrap_sequence: bool = True, + track_meta: Optional[bool] = None, allow_missing_keys: bool = False, ) -> None: """ @@ -533,28 +553,21 @@ def __init__( device: for Tensor data type, specify the target device. wrap_sequence: if `False`, then lists will recursively call this function, default to `True`. E.g., if `False`, `[1, 2]` -> `[tensor(1), tensor(2)]`, if `True`, then `[1, 2]` -> `tensor([1, 2])`. + track_meta: whether to convert to `MetaTensor` when `data_type` is "tensor". + If False, the output data type will be `torch.Tensor`. Default to the return value of `get_track_meta`. allow_missing_keys: don't raise exception if key is missing. """ super().__init__(keys, allow_missing_keys) - self.converter = EnsureType(data_type=data_type, dtype=dtype, device=device, wrap_sequence=wrap_sequence) + self.converter = EnsureType( + data_type=data_type, dtype=dtype, device=device, wrap_sequence=wrap_sequence, track_meta=track_meta + ) def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: d = dict(data) for key in self.key_iterator(d): - self.push_transform(d, key) d[key] = self.converter(d[key]) return d - def inverse(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]: - d = deepcopy(dict(data)) - for key in self.key_iterator(d): - # FIXME: currently, only convert tensor data to numpy array or scalar number, - # need to also invert numpy array but it's not easy to determine the previous data type - d[key] = convert_to_numpy(d[key]) - # Remove the applied transform - self.pop_transform(d, key) - return d - class ToNumpyd(MapTransform): """ @@ -975,6 +988,7 @@ class Lambdad(MapTransform, InvertibleTransform): print(lambd(input_data)['label'].shape) (4, 2, 2) + Args: keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` @@ -1007,29 +1021,20 @@ def __init__( self.overwrite = ensure_tuple_rep(overwrite, len(self.keys)) self._lambd = Lambda() - def _transform(self, data: Any, func: Callable): - return self._lambd(data, func=func) - - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: d = dict(data) for key, func, overwrite in self.key_iterator(d, self.func, self.overwrite): - ret = self._transform(data=d[key], func=func) + ret = self._lambd(img=d[key], func=func) if overwrite: d[key] = ret - self.push_transform(d, key) return d - def _inverse_transform(self, transform_info: Dict, data: Any, func: Callable): - return self._lambd(data, func=func) - def inverse(self, data): d = deepcopy(dict(data)) - for key, inv_func, overwrite in self.key_iterator(d, self.inv_func, self.overwrite): - transform = self.get_most_recent_transform(d, key) - ret = self._inverse_transform(transform_info=transform, data=d[key], func=inv_func) + for key, overwrite in self.key_iterator(d, self.overwrite): + ret = self._lambd.inverse(data=d[key]) if overwrite: d[key] = ret - self.pop_transform(d, key) return d @@ -1078,15 +1083,33 @@ def __init__( ) RandomizableTransform.__init__(self=self, prob=prob, do_transform=True) - def _transform(self, data: Any, func: Callable): - return self._lambd(data, func=func) if self._do_transform else data - def __call__(self, data): self.randomize(data) - return super().__call__(data) + d = dict(data) + for key, func, overwrite in self.key_iterator(d, self.func, self.overwrite): + ret = d[key] + if not isinstance(ret, MetaTensor): + ret = MetaTensor(ret) + if self._do_transform: + ret = self._lambd(ret, func=func) + self.push_transform(ret, extra_info={"lambda_info": self._lambd.pop_transform(ret)}) + else: + self.push_transform(ret) + if overwrite: + d[key] = ret + return d - def _inverse_transform(self, transform_info: Dict, data: Any, func: Callable): - return self._lambd(data, func=func) if transform_info[TraceKeys.DO_TRANSFORM] else data + def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]: + d = deepcopy(dict(data)) + for key, overwrite in self.key_iterator(d, self.overwrite): + if isinstance(d[key], MetaTensor): + tr = self.pop_transform(d[key]) + if tr[TraceKeys.DO_TRANSFORM]: + d[key].applied_operations.append(tr[TraceKeys.EXTRA_INFO]["lambda_info"]) # type: ignore + ret = self._lambd.inverse(d[key]) + if overwrite: + d[key] = ret + return d class LabelToMaskd(MapTransform): @@ -1222,7 +1245,7 @@ class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform): Dictionary-based wrapper of :py:class:`monai.transforms.ConvertToMultiChannelBasedOnBratsClasses`. Convert labels to multi channels based on brats18 classes: label 1 is the necrotic and non-enhancing tumor core - label 2 is the the peritumoral edema + label 2 is the peritumoral edema label 4 is the GD-enhancing tumor The possible classes are TC (Tumor core), WT (Whole tumor) and ET (Enhancing tumor). @@ -1420,7 +1443,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N class IntensityStatsd(MapTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.IntensityStats`. - Compute statistics for the intensity values of input image and store into the meta data dictionary. + Compute statistics for the intensity values of input image and store into the metadata dictionary. For example: if `ops=[lambda x: np.mean(x), "max"]` and `key_prefix="orig"`, may generate below stats: `{"orig_custom_0": 1.5, "orig_max": 3.0}`. @@ -1432,21 +1455,21 @@ class IntensityStatsd(MapTransform): mapping to `np.nanmean`, `np.nanmedian`, `np.nanmax`, `np.nanmin`, `np.nanstd`. if a callable function, will execute the function on input image. key_prefix: the prefix to combine with `ops` name to generate the key to store the results in the - meta data dictionary. if some `ops` are callable functions, will use "{key_prefix}_custom_{index}" + metadata dictionary. if some `ops` are callable functions, will use "{key_prefix}_custom_{index}" as the key, where index counts from 0. mask_keys: if not None, specify the mask array for the image to extract only the interested area to compute statistics, mask must have the same shape as the image. it should be a sequence of strings or None, map to the `keys`. channel_wise: whether to compute statistics for every channel of input image separately. if True, return a list of values for every operation, default to False. - meta_keys: explicitly indicate the key of the corresponding meta data dictionary. + meta_keys: explicitly indicate the key of the corresponding metadata dictionary. used to store the computed statistics to the meta dict. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. - the meta data is a dictionary object which contains: filename, original_shape, etc. + the metadata is a dictionary object which contains: filename, original_shape, etc. it can be a sequence of string, map to the `keys`. if None, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the meta data according - to the key data, default is `meta_dict`, the meta data is a dictionary object. + meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according + to the key data, default is `meta_dict`, the metadata is a dictionary object. used to store the computed statistics to the meta dict. allow_missing_keys: don't raise exception if key is missing. @@ -1637,6 +1660,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N RemoveRepeatedChannelD = RemoveRepeatedChannelDict = RemoveRepeatedChanneld RepeatChannelD = RepeatChannelDict = RepeatChanneld SplitChannelD = SplitChannelDict = SplitChanneld +SplitDimD = SplitDimDict = SplitDimd CastToTypeD = CastToTypeDict = CastToTyped ToTensorD = ToTensorDict = ToTensord EnsureTypeD = EnsureTypeDict = EnsureTyped diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 847614adfe..ccc467bda4 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -68,7 +68,7 @@ __all__ = [ "allow_missing_keys_mode", "compute_divisible_spatial_size", - "convert_inverse_interp_mode", + "convert_applied_interp_mode", "copypaste_arrays", "create_control_grid", "create_grid", @@ -105,6 +105,7 @@ "convert_pad_mode", "convert_to_contiguous", "get_unique_labels", + "scale_affine", ] @@ -576,7 +577,7 @@ def create_grid( dtype: Union[DtypeLike, torch.dtype] = float, device: Optional[torch.device] = None, backend=TransformBackends.NUMPY, -): +) -> NdarrayOrTensor: """ compute a `spatial_size` mesh. @@ -592,9 +593,9 @@ def create_grid( _backend = look_up_option(backend, TransformBackends) _dtype = dtype or float if _backend == TransformBackends.NUMPY: - return _create_grid_numpy(spatial_size, spacing, homogeneous, _dtype) + return _create_grid_numpy(spatial_size, spacing, homogeneous, _dtype) # type: ignore if _backend == TransformBackends.TORCH: - return _create_grid_torch(spatial_size, spacing, homogeneous, _dtype, device) + return _create_grid_torch(spatial_size, spacing, homogeneous, _dtype, device) # type: ignore raise ValueError(f"backend {backend} is not supported") @@ -671,7 +672,7 @@ def create_rotate( spatial_dims: int, radians: Union[Sequence[float], float], device: Optional[torch.device] = None, - backend=TransformBackends.NUMPY, + backend: str = TransformBackends.NUMPY, ) -> NdarrayOrTensor: """ create a 2D or 3D rotation matrix @@ -934,8 +935,8 @@ def generate_spatial_bounding_box( min_d = max(min_d, 0) max_d = min(max_d, spatial_size[di]) - box_start[di] = min_d.detach().cpu().item() if isinstance(min_d, torch.Tensor) else min_d # type: ignore - box_end[di] = max_d.detach().cpu().item() if isinstance(max_d, torch.Tensor) else max_d # type: ignore + box_start[di] = min_d.detach().cpu().item() if isinstance(min_d, torch.Tensor) else min_d + box_end[di] = max_d.detach().cpu().item() if isinstance(max_d, torch.Tensor) else max_d return box_start, box_end @@ -1185,16 +1186,13 @@ def map_spatial_axes( """ if spatial_axes is None: - spatial_axes_ = list(range(1, img_ndim) if channel_first else range(img_ndim - 1)) - - else: - spatial_axes_ = [] - for a in ensure_tuple(spatial_axes): - if channel_first: - spatial_axes_.append(a if a < 0 else a + 1) - else: - spatial_axes_.append(a - 1 if a < 0 else a) - + return list(range(1, img_ndim) if channel_first else range(img_ndim - 1)) + spatial_axes_ = [] + for a in ensure_tuple(spatial_axes): + if channel_first: + spatial_axes_.append(a % img_ndim if a < 0 else a + 1) + else: + spatial_axes_.append((a - 1) % (img_ndim - 1) if a < 0 else a) return spatial_axes_ @@ -1245,37 +1243,42 @@ def allow_missing_keys_mode(transform: Union[MapTransform, Compose, Tuple[MapTra t.allow_missing_keys = o_s -def convert_inverse_interp_mode(trans_info: List, mode: str = "nearest", align_corners: Optional[bool] = None): +_interp_modes = list(InterpolateMode) + list(GridSampleMode) + + +def convert_applied_interp_mode(trans_info, mode: str = "nearest", align_corners: Optional[bool] = None): """ - Change the interpolation mode when inverting spatial transforms, default to "nearest". - This function modifies trans_info's `TraceKeys.EXTRA_INFO`. + Recursively change the interpolation mode in the applied operation stacks, default to "nearest". See also: :py:class:`monai.transform.inverse.InvertibleTransform` Args: - trans_info: transforms inverse information list, contains context of every invertible transform. + trans_info: applied operation stack, tracking the previously applied invertible transform. mode: target interpolation mode to convert, default to "nearest" as it's usually used to save the mode output. align_corners: target align corner value in PyTorch interpolation API, need to align with the `mode`. """ - interp_modes = [i.value for i in InterpolateMode] + [i.value for i in GridSampleMode] - - # set to string for DataLoader collation - align_corners_ = TraceKeys.NONE if align_corners is None else align_corners - - for item in ensure_tuple(trans_info): - if TraceKeys.EXTRA_INFO in item: - orig_mode = item[TraceKeys.EXTRA_INFO].get("mode", None) - if orig_mode is not None: - if orig_mode[0] in interp_modes: - item[TraceKeys.EXTRA_INFO]["mode"] = [mode for _ in range(len(mode))] - elif orig_mode in interp_modes: - item[TraceKeys.EXTRA_INFO]["mode"] = mode - if "align_corners" in item[TraceKeys.EXTRA_INFO]: - if issequenceiterable(item[TraceKeys.EXTRA_INFO]["align_corners"]): - item[TraceKeys.EXTRA_INFO]["align_corners"] = [align_corners_ for _ in range(len(mode))] - else: - item[TraceKeys.EXTRA_INFO]["align_corners"] = align_corners_ + if isinstance(trans_info, (list, tuple)): + return [convert_applied_interp_mode(x, mode=mode, align_corners=align_corners) for x in trans_info] + if not isinstance(trans_info, Mapping): + return trans_info + trans_info = dict(trans_info) + if "mode" in trans_info: + current_mode = trans_info["mode"] + if current_mode[0] in _interp_modes: + trans_info["mode"] = [mode for _ in range(len(mode))] + elif current_mode in _interp_modes: + trans_info["mode"] = mode + if "align_corners" in trans_info: + _align_corners = TraceKeys.NONE if align_corners is None else align_corners + current_value = trans_info["align_corners"] + trans_info["align_corners"] = ( + [_align_corners for _ in mode] if issequenceiterable(current_value) else _align_corners + ) + if ("mode" not in trans_info) and ("align_corners" not in trans_info): + return { + k: convert_applied_interp_mode(trans_info[k], mode=mode, align_corners=align_corners) for k in trans_info + } return trans_info @@ -1529,7 +1532,7 @@ def print_table_column(name, torch, numpy, color=Colors.none): print_color(f"Number of uncategorised: {n_uncategorized}", Colors.red) -def convert_pad_mode(dst: NdarrayOrTensor, mode: Union[NumpyPadMode, PytorchPadMode, str]): +def convert_pad_mode(dst: NdarrayOrTensor, mode: Optional[str]): """ Utility to convert padding mode between numpy array and PyTorch Tensor. @@ -1538,7 +1541,6 @@ def convert_pad_mode(dst: NdarrayOrTensor, mode: Union[NumpyPadMode, PytorchPadM mode: current padding mode. """ - mode = mode.value if isinstance(mode, (NumpyPadMode, PytorchPadMode)) else mode if isinstance(dst, torch.Tensor): if mode == "wrap": mode = "circular" @@ -1556,7 +1558,7 @@ def convert_pad_mode(dst: NdarrayOrTensor, mode: Union[NumpyPadMode, PytorchPadM def convert_to_contiguous(data, **kwargs): """ - Check and ensure the numpy array or PyTorch Tensor in data to be contuguous in memory. + Check and ensure the numpy array or PyTorch Tensor in data to be contiguous in memory. Args: data: input data to convert, will recursively convert the numpy array or PyTorch Tensor in dict and sequence. @@ -1573,5 +1575,30 @@ def convert_to_contiguous(data, **kwargs): return data +def scale_affine(affine, spatial_size, new_spatial_size, centered: bool = True): + """ + Scale the affine matrix according to the new spatial size. + + Args: + affine: affine matrix to scale. + spatial_size: original spatial size. + new_spatial_size: new spatial size. + centered: whether the scaling is with respect to + the image center (True, default) or corner (False). + + Returns: + Scaled affine matrix. + + """ + if spatial_size == new_spatial_size: + return affine + r = len(affine) - 1 + s = np.array([float(o) / float(max(n, 1)) for o, n in zip(spatial_size, new_spatial_size)]) + scale = create_scale(r, s.tolist()) + if centered: + scale[:r, -1] = (np.diag(scale)[:r] - 1) / 2 # type: ignore + return affine @ convert_to_dst_type(scale, affine)[0] + + if __name__ == "__main__": print_transform_backends() diff --git a/monai/transforms/utils_create_transform_ims.py b/monai/transforms/utils_create_transform_ims.py index b096e1b93d..8f2ae82639 100644 --- a/monai/transforms/utils_create_transform_ims.py +++ b/monai/transforms/utils_create_transform_ims.py @@ -85,6 +85,7 @@ ) from monai.transforms.intensity.array import ( AdjustContrast, + ForegroundMask, GaussianSharpen, GaussianSmooth, GibbsNoise, @@ -115,6 +116,7 @@ ) from monai.transforms.intensity.dictionary import ( AdjustContrastd, + ForegroundMaskd, GaussianSharpend, GaussianSmoothd, GibbsNoised, @@ -458,7 +460,6 @@ def create_transform_im( create_transform_im(RandFlipd, dict(keys=keys, prob=1, spatial_axis=2), data) create_transform_im(Flip, dict(spatial_axis=1), data) create_transform_im(Flipd, dict(keys=keys, spatial_axis=2), data) - create_transform_im(Flipd, dict(keys=keys, spatial_axis=2), data) create_transform_im(Orientation, dict(axcodes="RPI", image_only=True), data) create_transform_im(Orientationd, dict(keys=keys, axcodes="RPI"), data) create_transform_im( @@ -584,6 +585,8 @@ def create_transform_im( create_transform_im( MaskIntensityd, dict(keys=CommonKeys.IMAGE, mask_key=CommonKeys.IMAGE, select_fn=lambda x: x > 0.3), data ) + create_transform_im(ForegroundMask, dict(invert=True), data) + create_transform_im(ForegroundMaskd, dict(keys=CommonKeys.IMAGE, invert=True), data) create_transform_im(GaussianSmooth, dict(sigma=2), data) create_transform_im(GaussianSmoothd, dict(keys=CommonKeys.IMAGE, sigma=2), data) create_transform_im(RandGaussianSmooth, dict(prob=1.0, sigma_x=(1, 2)), data) @@ -718,7 +721,7 @@ def create_transform_im( create_transform_im( RandSmoothDeform, - dict(spatial_size=(217, 217, 217), rand_size=(10, 10, 10), prob=1.0, def_range=0.05, grid_mode="blinear"), + dict(spatial_size=(217, 217, 217), rand_size=(10, 10, 10), prob=1.0, def_range=0.05, grid_mode="bilinear"), data, ) create_transform_im( @@ -729,7 +732,7 @@ def create_transform_im( rand_size=(10, 10, 10), prob=1.0, def_range=0.05, - grid_mode="blinear", + grid_mode="bilinear", ), data, ) diff --git a/monai/transforms/utils_pytorch_numpy_unification.py b/monai/transforms/utils_pytorch_numpy_unification.py index 2103ccff58..441ea23b2f 100644 --- a/monai/transforms/utils_pytorch_numpy_unification.py +++ b/monai/transforms/utils_pytorch_numpy_unification.py @@ -15,7 +15,7 @@ import torch from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor -from monai.utils.misc import ensure_tuple, is_module_ver_at_least +from monai.utils.misc import is_module_ver_at_least from monai.utils.type_conversion import convert_data_type, convert_to_dst_type __all__ = [ @@ -54,31 +54,12 @@ def allclose(a: NdarrayTensor, b: NdarrayOrTensor, rtol=1e-5, atol=1e-8, equal_n def moveaxis(x: NdarrayOrTensor, src: Union[int, Sequence[int]], dst: Union[int, Sequence[int]]) -> NdarrayOrTensor: - """`moveaxis` for pytorch and numpy, using `permute` for pytorch version < 1.7""" + """`moveaxis` for pytorch and numpy""" if isinstance(x, torch.Tensor): - if hasattr(torch, "movedim"): # `movedim` is new in torch 1.7.0 - # torch.moveaxis is a recent alias since torch 1.8.0 - return torch.movedim(x, src, dst) # type: ignore - return _moveaxis_with_permute(x, src, dst) + return torch.movedim(x, src, dst) # type: ignore return np.moveaxis(x, src, dst) -def _moveaxis_with_permute( - x: torch.Tensor, src: Union[int, Sequence[int]], dst: Union[int, Sequence[int]] -) -> torch.Tensor: - # get original indices - indices = list(range(x.ndim)) - len_indices = len(indices) - for s, d in zip(ensure_tuple(src), ensure_tuple(dst)): - # make src and dst positive - # remove desired index and insert it in new position - pos_s = len_indices + s if s < 0 else s - pos_d = len_indices + d if d < 0 else d - indices.pop(pos_s) - indices.insert(pos_d, pos_s) - return x.permute(indices) - - def in1d(x, y): """`np.in1d` with equivalent implementation for torch.""" if isinstance(x, np.ndarray): @@ -101,10 +82,7 @@ def percentile( ) -> Union[NdarrayOrTensor, float, int]: """`np.percentile` with equivalent implementation for torch. - Pytorch uses `quantile`, but this functionality is only available from v1.7. - For earlier methods, we calculate it ourselves. This doesn't do interpolation, - so is the equivalent of ``numpy.percentile(..., interpolation="nearest")``. - For more details, please refer to: + Pytorch uses `quantile`. For more details please refer to: https://pytorch.org/docs/stable/generated/torch.quantile.html. https://numpy.org/doc/stable/reference/generated/numpy.percentile.html. @@ -112,7 +90,7 @@ def percentile( x: input data q: percentile to compute (should in range 0 <= q <= 100) dim: the dim along which the percentiles are computed. default is to compute the percentile - along a flattened version of the array. only work for numpy array or Tensor with PyTorch >= 1.7.0. + along a flattened version of the array. keepdim: whether the output data has dim retained or not. kwargs: if `x` is numpy array, additional args for `np.percentile`, more details: https://numpy.org/doc/stable/reference/generated/numpy.percentile.html. @@ -130,18 +108,7 @@ def percentile( result = np.percentile(x, q, axis=dim, keepdims=keepdim, **kwargs) else: q = torch.tensor(q, device=x.device) - if hasattr(torch, "quantile"): # `quantile` is new in torch 1.7.0 - result = torch.quantile(x, q / 100.0, dim=dim, keepdim=keepdim) - else: - # Note that ``kthvalue()`` works one-based, i.e., the first sorted value - # corresponds to k=1, not k=0. Thus, we need the `1 +`. - k = 1 + (0.01 * q * (x.numel() - 1)).round().int() - if k.numel() > 1: - r = [x.view(-1).kthvalue(int(_k)).values.item() for _k in k] - result = torch.tensor(r, device=x.device) - else: - result = x.view(-1).kthvalue(int(k)).values.item() - + result = torch.quantile(x, q / 100.0, dim=dim, keepdim=keepdim) return result @@ -277,8 +244,6 @@ def any_np_pt(x: NdarrayOrTensor, axis: Union[int, Sequence[int]]) -> NdarrayOrT def maximum(a: NdarrayOrTensor, b: NdarrayOrTensor) -> NdarrayOrTensor: """`np.maximum` with equivalent implementation for torch. - `torch.maximum` only available from pt>1.6, else use `torch.stack` and `torch.max`. - Args: a: first array/tensor b: second array/tensor @@ -287,10 +252,7 @@ def maximum(a: NdarrayOrTensor, b: NdarrayOrTensor) -> NdarrayOrTensor: Element-wise maximum between two arrays/tensors. """ if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): - # is torch and has torch.maximum (pt>1.6) - if hasattr(torch, "maximum"): # `maximum` is new in torch 1.7.0 - return torch.maximum(a, b) - return torch.stack((a, b)).max(dim=0)[0] + return torch.maximum(a, b) return np.maximum(a, b) @@ -314,7 +276,7 @@ def cumsum(a: NdarrayOrTensor, axis=None, **kwargs) -> NdarrayOrTensor: """ if isinstance(a, np.ndarray): - return np.cumsum(a, axis) + return np.cumsum(a, axis) # type: ignore if axis is None: return torch.cumsum(a[:], 0, **kwargs) return torch.cumsum(a, dim=axis, **kwargs) @@ -352,7 +314,7 @@ def repeat(a: NdarrayOrTensor, repeats: int, axis: Optional[int] = None, **kwarg Args: a: input data to repeat. - repeats: number of repetitions for each element, repeats is broadcasted to fit the shape of the given axis. + repeats: number of repetitions for each element, repeats is broadcast to fit the shape of the given axis. axis: axis along which to repeat values. kwargs: if `a` is PyTorch Tensor, additional args for `torch.repeat_interleave`, more details: https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html. @@ -427,3 +389,14 @@ def unique(x: NdarrayTensor) -> NdarrayTensor: x: array/tensor """ return torch.unique(x) if isinstance(x, torch.Tensor) else np.unique(x) # type: ignore + + +def linalg_inv(x: NdarrayTensor) -> NdarrayTensor: + """`torch.linalg.inv` with equivalent implementation for numpy. + + Args: + x: array/tensor + """ + if isinstance(x, torch.Tensor) and hasattr(torch, "inverse"): # pytorch 1.7.0 + return torch.inverse(x) # type: ignore + return torch.linalg.inv(x) if isinstance(x, torch.Tensor) else np.linalg.inv(x) # type: ignore diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index 429183b1a0..cb376b7280 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -17,10 +17,12 @@ from .enums import ( Average, BlendMode, + BoxModeName, ChannelMatching, CommonKeys, DiceCEReduction, ForwardMode, + GridPatchSort, GridSampleMode, GridSamplePadMode, InterpolateMode, @@ -31,12 +33,15 @@ MetricReduction, NumpyPadMode, PostFix, + ProbMapKeys, PytorchPadMode, SkipMode, + StrEnum, TraceKeys, TransformBackends, UpsampleMode, Weight, + WSIPatchKeys, ) from .jupyter_utils import StatusMembers, ThreadContainer from .misc import ( @@ -88,10 +93,13 @@ convert_data_type, convert_to_cupy, convert_to_dst_type, + convert_to_list, convert_to_numpy, convert_to_tensor, dtype_numpy_to_torch, dtype_torch_to_numpy, get_dtype, get_equivalent_dtype, + get_numpy_dtype_from_string, + get_torch_dtype_from_string, ) diff --git a/monai/utils/enums.py b/monai/utils/enums.py index 4bc3d6ee84..88faf88432 100644 --- a/monai/utils/enums.py +++ b/monai/utils/enums.py @@ -9,12 +9,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import random from enum import Enum from typing import Optional from monai.utils.deprecate_utils import deprecated __all__ = [ + "StrEnum", "NumpyPadMode", "GridSampleMode", "InterpolateMode", @@ -36,10 +38,37 @@ "PostFix", "ForwardMode", "TransformBackends", + "BoxModeName", + "GridPatchSort", ] -class NumpyPadMode(Enum): +class StrEnum(str, Enum): + """ + Enum subclass that converts its value to a string. + + .. code-block:: python + + from monai.utils import StrEnum + + class Example(StrEnum): + MODE_A = "A" + MODE_B = "B" + + assert (list(Example) == ["A", "B"]) + assert Example.MODE_A == "A" + assert str(Example.MODE_A) == "A" + assert monai.utils.look_up_option("A", Example) == "A" + """ + + def __str__(self): + return self.value + + def __repr__(self): + return self.value + + +class NumpyPadMode(StrEnum): """ See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html """ @@ -57,7 +86,7 @@ class NumpyPadMode(Enum): EMPTY = "empty" -class GridSampleMode(Enum): +class GridSampleMode(StrEnum): """ See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html @@ -75,12 +104,13 @@ class GridSampleMode(Enum): BICUBIC = "bicubic" -class InterpolateMode(Enum): +class InterpolateMode(StrEnum): """ See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html """ NEAREST = "nearest" + NEAREST_EXACT = "nearest-exact" LINEAR = "linear" BILINEAR = "bilinear" BICUBIC = "bicubic" @@ -88,7 +118,7 @@ class InterpolateMode(Enum): AREA = "area" -class UpsampleMode(Enum): +class UpsampleMode(StrEnum): """ See also: :py:class:`monai.networks.blocks.UpSample` """ @@ -98,7 +128,7 @@ class UpsampleMode(Enum): PIXELSHUFFLE = "pixelshuffle" -class BlendMode(Enum): +class BlendMode(StrEnum): """ See also: :py:class:`monai.data.utils.compute_importance_map` """ @@ -107,7 +137,7 @@ class BlendMode(Enum): GAUSSIAN = "gaussian" -class PytorchPadMode(Enum): +class PytorchPadMode(StrEnum): """ See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html """ @@ -118,7 +148,7 @@ class PytorchPadMode(Enum): CIRCULAR = "circular" -class GridSamplePadMode(Enum): +class GridSamplePadMode(StrEnum): """ See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html """ @@ -128,7 +158,7 @@ class GridSamplePadMode(Enum): REFLECTION = "reflection" -class Average(Enum): +class Average(StrEnum): """ See also: :py:class:`monai.metrics.rocauc.compute_roc_auc` """ @@ -139,9 +169,9 @@ class Average(Enum): NONE = "none" -class MetricReduction(Enum): +class MetricReduction(StrEnum): """ - See also: :py:class:`monai.metrics.meandice.DiceMetric` + See also: :py:func:`monai.metrics.utils.do_metric_reduction` """ NONE = "none" @@ -153,7 +183,7 @@ class MetricReduction(Enum): SUM_CHANNEL = "sum_channel" -class LossReduction(Enum): +class LossReduction(StrEnum): """ See also: - :py:class:`monai.losses.dice.DiceLoss` @@ -167,7 +197,7 @@ class LossReduction(Enum): SUM = "sum" -class DiceCEReduction(Enum): +class DiceCEReduction(StrEnum): """ See also: - :py:class:`monai.losses.dice.DiceCELoss` @@ -177,7 +207,7 @@ class DiceCEReduction(Enum): SUM = "sum" -class Weight(Enum): +class Weight(StrEnum): """ See also: :py:class:`monai.losses.dice.GeneralizedDiceLoss` """ @@ -187,7 +217,7 @@ class Weight(Enum): UNIFORM = "uniform" -class ChannelMatching(Enum): +class ChannelMatching(StrEnum): """ See also: :py:class:`monai.networks.nets.HighResBlock` """ @@ -196,7 +226,7 @@ class ChannelMatching(Enum): PROJECT = "project" -class SkipMode(Enum): +class SkipMode(StrEnum): """ See also: :py:class:`monai.networks.layers.SkipConnection` """ @@ -206,7 +236,7 @@ class SkipMode(Enum): MUL = "mul" -class Method(Enum): +class Method(StrEnum): """ See also: :py:class:`monai.transforms.croppad.array.SpatialPad` """ @@ -215,7 +245,7 @@ class Method(Enum): END = "end" -class ForwardMode(Enum): +class ForwardMode(StrEnum): """ See also: :py:class:`monai.transforms.engines.evaluator.Evaluator` """ @@ -224,22 +254,22 @@ class ForwardMode(Enum): EVAL = "eval" -class TraceKeys: - """Extra meta data keys used for traceable transforms.""" +class TraceKeys(StrEnum): + """Extra metadata keys used for traceable transforms.""" - CLASS_NAME = "class" - ID = "id" - ORIG_SIZE = "orig_size" - EXTRA_INFO = "extra_info" - DO_TRANSFORM = "do_transforms" - KEY_SUFFIX = "_transforms" - NONE = "none" + CLASS_NAME: str = "class" + ID: str = "id" + ORIG_SIZE: str = "orig_size" + EXTRA_INFO: str = "extra_info" + DO_TRANSFORM: str = "do_transforms" + KEY_SUFFIX: str = "_transforms" + NONE: str = "none" @deprecated(since="0.8.0", msg_suffix="use monai.utils.enums.TraceKeys instead.") class InverseKeys: """ - Extra meta data keys used for inverse transforms. + Extra metadata keys used for inverse transforms. .. deprecated:: 0.8.0 Use :class:`monai.utils.enums.TraceKeys` instead. @@ -255,7 +285,7 @@ class InverseKeys: NONE = "none" -class CommonKeys: +class CommonKeys(StrEnum): """ A set of common keys for dictionary based supervised training process. `IMAGE` is the input image data. @@ -270,9 +300,10 @@ class CommonKeys: LABEL = "label" PRED = "pred" LOSS = "loss" + METADATA = "metadata" -class PostFix: +class PostFix(StrEnum): """Post-fixes.""" @staticmethod @@ -287,8 +318,12 @@ def meta(key: Optional[str] = None): def orig_meta(key: Optional[str] = None): return PostFix._get_str(key, "orig_meta_dict") + @staticmethod + def transforms(key: Optional[str] = None): + return PostFix._get_str(key, TraceKeys.KEY_SUFFIX[1:]) + -class TransformBackends(Enum): +class TransformBackends(StrEnum): """ Transform backends. """ @@ -297,7 +332,7 @@ class TransformBackends(Enum): NUMPY = "numpy" -class JITMetadataKeys(Enum): +class JITMetadataKeys(StrEnum): """ Keys stored in the metadata file for saved Torchscript models. Some of these are generated by the routines and others are optionally provided by users. @@ -307,3 +342,74 @@ class JITMetadataKeys(Enum): TIMESTAMP = "timestamp" VERSION = "version" DESCRIPTION = "description" + + +class BoxModeName(StrEnum): + """ + Box mode names. + """ + + XYXY = "xyxy" # [xmin, ymin, xmax, ymax] + XYZXYZ = "xyzxyz" # [xmin, ymin, zmin, xmax, ymax, zmax] + XXYY = "xxyy" # [xmin, xmax, ymin, ymax] + XXYYZZ = "xxyyzz" # [xmin, xmax, ymin, ymax, zmin, zmax] + XYXYZZ = "xyxyzz" # [xmin, ymin, xmax, ymax, zmin, zmax] + XYWH = "xywh" # [xmin, ymin, xsize, ysize] + XYZWHD = "xyzwhd" # [xmin, ymin, zmin, xsize, ysize, zsize] + CCWH = "ccwh" # [xcenter, ycenter, xsize, ysize] + CCCWHD = "cccwhd" # [xcenter, ycenter, zcenter, xsize, ysize, zsize] + + +class ProbMapKeys(StrEnum): + """ + The keys to be used for generating the probability maps from patches + """ + + LOCATION = "mask_location" + SIZE = "mask_size" + COUNT = "num_patches" + NAME = "name" + + +class GridPatchSort(StrEnum): + """ + The sorting method for the generated patches in `GridPatch` + """ + + RANDOM = "random" + MIN = "min" + MAX = "max" + + @staticmethod + def min_fn(x): + return x[0].sum() + + @staticmethod + def max_fn(x): + return -x[0].sum() + + @staticmethod + def get_sort_fn(sort_fn): + if sort_fn == GridPatchSort.RANDOM: + return random.random + elif sort_fn == GridPatchSort.MIN: + return GridPatchSort.min_fn + elif sort_fn == GridPatchSort.MAX: + return GridPatchSort.max_fn + else: + raise ValueError( + f'sort_fn should be one of the following values, "{sort_fn}" was given:', + [e.value for e in GridPatchSort], + ) + + +class WSIPatchKeys(StrEnum): + """ + The keys to be used for metadata of patches extracted from whole slide images + """ + + LOCATION = "patch_location" + LEVEL = "patch_level" + SIZE = "patch_size" + COUNT = "num_patches" + PATH = "path" diff --git a/monai/utils/misc.py b/monai/utils/misc.py index 5f074289bc..fc38dc5056 100644 --- a/monai/utils/misc.py +++ b/monai/utils/misc.py @@ -9,7 +9,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import collections.abc import inspect import itertools import os @@ -19,6 +18,7 @@ import types import warnings from ast import literal_eval +from collections.abc import Iterable from distutils.util import strtobool from pathlib import Path from typing import Any, Callable, Optional, Sequence, Tuple, Union, cast @@ -26,7 +26,7 @@ import numpy as np import torch -from monai.config.type_definitions import NdarrayOrTensor, PathLike +from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor, PathLike from monai.utils.module import version_leq __all__ = [ @@ -88,19 +88,27 @@ def issequenceiterable(obj: Any) -> bool: """ Determine if the object is an iterable sequence and is not a string. """ - if isinstance(obj, torch.Tensor): - return int(obj.dim()) > 0 # a 0-d tensor is not iterable - return isinstance(obj, collections.abc.Iterable) and not isinstance(obj, (str, bytes)) + try: + if hasattr(obj, "ndim") and obj.ndim == 0: + return False # a 0-d tensor is not iterable + except Exception: + return False + return isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)) -def ensure_tuple(vals: Any) -> Tuple[Any, ...]: +def ensure_tuple(vals: Any, wrap_array: bool = False) -> Tuple[Any, ...]: """ Returns a tuple of `vals`. + + Args: + vals: input data to convert to a tuple. + wrap_array: if `True`, treat the input numerical array (ndarray/tensor) as one item of the tuple. + if `False`, try to convert the array with `tuple(vals)`, default to `False`. + """ - if not issequenceiterable(vals): + if wrap_array and isinstance(vals, (np.ndarray, torch.Tensor)): return (vals,) - - return tuple(vals) + return tuple(vals) if issequenceiterable(vals) else (vals,) def ensure_tuple_size(tup: Any, dim: int, pad_val: Any = 0) -> Tuple[Any, ...]: @@ -147,7 +155,7 @@ def ensure_tuple_rep(tup: Any, dim: int) -> Tuple[Any, ...]: def fall_back_tuple( - user_provided: Any, default: Union[Sequence, np.ndarray], func: Callable = lambda x: x and x > 0 + user_provided: Any, default: Union[Sequence, NdarrayTensor], func: Callable = lambda x: x and x > 0 ) -> Tuple[Any, ...]: """ Refine `user_provided` according to the `default`, and returns as a validated tuple. @@ -354,11 +362,12 @@ def copy_to_device( class ImageMetaKey: """ - Common key names in the meta data header of images + Common key names in the metadata header of images """ FILENAME_OR_OBJ = "filename_or_obj" PATCH_INDEX = "patch_index" + SPATIAL_SHAPE = "spatial_shape" def has_option(obj, keywords: Union[str, Sequence[str]]) -> bool: diff --git a/monai/utils/module.py b/monai/utils/module.py index 065cc8f7c8..747c985af7 100644 --- a/monai/utils/module.py +++ b/monai/utils/module.py @@ -213,11 +213,14 @@ def instantiate(path: str, **kwargs): component = locate(path) if component is None: raise ModuleNotFoundError(f"Cannot locate class or function path: '{path}'.") - if isclass(component): - return component(**kwargs) - # support regular function, static method and class method - if isfunction(component) or (ismethod(component) and isclass(getattr(component, "__self__", None))): - return partial(component, **kwargs) + try: + if isclass(component): + return component(**kwargs) + # support regular function, static method and class method + if isfunction(component) or (ismethod(component) and isclass(getattr(component, "__self__", None))): + return partial(component, **kwargs) + except Exception as e: + raise RuntimeError(f"Failed to instantiate '{path}' with kwargs: {kwargs}") from e warnings.warn(f"Component to instantiate must represent a valid class or function, but got {path}.") return component diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 691f900c7d..fefab380f1 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -43,6 +43,8 @@ class Range: Otherwise, it look up predefined methods: "forward", "__call__", "__next__", "__getitem__" append_method_name: if append the name of the methods to be decorated to the range's name If None (default), it appends the method's name only if we are annotating more than one method. + recursive: if set to True, it will recursively annotate every individual module in a list + or in a chain of modules (chained using Compose). Default to False. """ @@ -53,12 +55,25 @@ def __init__( name: Optional[str] = None, methods: Optional[Union[str, Tuple[str, ...]]] = None, append_method_name: Optional[bool] = None, + recursive: bool = False, ) -> None: self.name = name self.methods = methods self.append_method_name = append_method_name + self.recursive = recursive def __call__(self, obj: Any): + if self.recursive is True: + if isinstance(obj, (list, tuple)): + return type(obj)(Range(recursive=True)(t) for t in obj) + + from monai.transforms.compose import Compose + + if isinstance(obj, Compose): + obj.transforms = Range(recursive=True)(obj.transforms) + + self.recursive = False + # Define the name to be associated to the range if not provided if self.name is None: name = type(obj).__name__ diff --git a/monai/utils/type_conversion.py b/monai/utils/type_conversion.py index d5944e265b..e33a155568 100644 --- a/monai/utils/type_conversion.py +++ b/monai/utils/type_conversion.py @@ -10,19 +10,22 @@ # limitations under the License. import re +from copy import deepcopy from typing import Any, Optional, Sequence, Tuple, Type, Union import numpy as np import torch +import monai from monai.config.type_definitions import DtypeLike, NdarrayTensor from monai.utils import optional_import -from monai.utils.module import look_up_option cp, has_cp = optional_import("cupy") cp_ndarray, _ = optional_import("cupy", name="ndarray") __all__ = [ + "get_numpy_dtype_from_string", + "get_torch_dtype_from_string", "dtype_torch_to_numpy", "dtype_numpy_to_torch", "get_equivalent_dtype", @@ -35,37 +38,32 @@ ] -_torch_to_np_dtype = { - torch.bool: np.dtype(bool), - torch.uint8: np.dtype(np.uint8), - torch.int8: np.dtype(np.int8), - torch.int16: np.dtype(np.int16), - torch.int32: np.dtype(np.int32), - torch.int64: np.dtype(np.int64), - torch.float16: np.dtype(np.float16), - torch.float32: np.dtype(np.float32), - torch.float64: np.dtype(np.float64), - torch.complex64: np.dtype(np.complex64), - torch.complex128: np.dtype(np.complex128), -} -_np_to_torch_dtype = {value: key for key, value in _torch_to_np_dtype.items()} +def get_numpy_dtype_from_string(dtype: str) -> np.dtype: + """Get a numpy dtype (e.g., `np.float32`) from its string (e.g., `"float32"`).""" + return np.empty([], dtype=dtype).dtype -def dtype_torch_to_numpy(dtype): +def get_torch_dtype_from_string(dtype: str) -> torch.dtype: + """Get a torch dtype (e.g., `torch.float32`) from its string (e.g., `"float32"`).""" + return dtype_numpy_to_torch(get_numpy_dtype_from_string(dtype)) + + +def dtype_torch_to_numpy(dtype: torch.dtype) -> np.dtype: """Convert a torch dtype to its numpy equivalent.""" - return look_up_option(dtype, _torch_to_np_dtype) + return torch.empty([], dtype=dtype).numpy().dtype # type: ignore -def dtype_numpy_to_torch(dtype): +def dtype_numpy_to_torch(dtype: np.dtype) -> torch.dtype: """Convert a numpy dtype to its torch equivalent.""" - # np dtypes can be given as np.float32 and np.dtype(np.float32) so unify them - dtype = np.dtype(dtype) if isinstance(dtype, (type, str)) else dtype - return look_up_option(dtype, _np_to_torch_dtype) + return torch.from_numpy(np.empty([], dtype=dtype)).dtype def get_equivalent_dtype(dtype, data_type): """Convert to the `dtype` that corresponds to `data_type`. + The input dtype can also be a string. e.g., `"float32"` becomes `torch.float32` or + `np.float32` as necessary. + Example:: im = torch.tensor(1) @@ -74,7 +72,7 @@ def get_equivalent_dtype(dtype, data_type): """ if dtype is None: return None - if data_type is torch.Tensor: + if data_type is torch.Tensor or data_type.__name__ == "MetaTensor": if isinstance(dtype, torch.dtype): # already a torch dtype and target `data_type` is torch.Tensor return dtype @@ -100,11 +98,16 @@ def get_dtype(data: Any): def convert_to_tensor( - data, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, wrap_sequence: bool = False + data, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + wrap_sequence: bool = False, + track_meta: bool = False, ): """ - Utility to convert the input data to a PyTorch Tensor. If passing a dictionary, list or tuple, - recursively check every item and convert it to PyTorch Tensor. + Utility to convert the input data to a PyTorch Tensor, if `track_meta` is True, the output will be a `MetaTensor`, + otherwise, the output will be a regular torch Tensor. + If passing a dictionary, list or tuple, recursively check every item and convert it to PyTorch Tensor. Args: data: input data can be PyTorch Tensor, numpy array, list, dictionary, int, float, bool, str, etc. @@ -114,10 +117,24 @@ def convert_to_tensor( device: target device to put the converted Tensor data. wrap_sequence: if `False`, then lists will recursively call this function. E.g., `[1, 2]` -> `[tensor(1), tensor(2)]`. If `True`, then `[1, 2]` -> `tensor([1, 2])`. + track_meta: whether to track the meta information, if `True`, will convert to `MetaTensor`. + default to `False`. """ + + def _convert_tensor(tensor, **kwargs): + if not isinstance(tensor, torch.Tensor): + # if input data is not Tensor, convert it to Tensor first + tensor = torch.as_tensor(tensor, **kwargs) + if track_meta and not isinstance(tensor, monai.data.MetaTensor): + return monai.data.MetaTensor(tensor) + if not track_meta and isinstance(tensor, monai.data.MetaTensor): + return tensor.as_tensor() + return tensor + + dtype = get_equivalent_dtype(dtype, torch.Tensor) if isinstance(data, torch.Tensor): - return data.to(dtype=dtype, device=device, memory_format=torch.contiguous_format) # type: ignore + return _convert_tensor(data).to(dtype=dtype, device=device, memory_format=torch.contiguous_format) if isinstance(data, np.ndarray): # skip array of string classes and object, refer to: # https://github.com/pytorch/pytorch/blob/v1.9.0/torch/utils/data/_utils/collate.py#L13 @@ -126,17 +143,17 @@ def convert_to_tensor( # `ascontiguousarray` will add 1 dim if img has no dim, so we only apply on data with dims if data.ndim > 0: data = np.ascontiguousarray(data) - return torch.as_tensor(data, dtype=dtype, device=device) # type: ignore + return _convert_tensor(data, dtype=dtype, device=device) elif (has_cp and isinstance(data, cp_ndarray)) or isinstance(data, (float, int, bool)): - return torch.as_tensor(data, dtype=dtype, device=device) # type: ignore + return _convert_tensor(data, dtype=dtype, device=device) elif isinstance(data, list): - list_ret = [convert_to_tensor(i, dtype=dtype, device=device) for i in data] - return torch.as_tensor(list_ret, dtype=dtype, device=device) if wrap_sequence else list_ret # type: ignore + list_ret = [convert_to_tensor(i, dtype=dtype, device=device, track_meta=track_meta) for i in data] + return _convert_tensor(list_ret, dtype=dtype, device=device) if wrap_sequence else list_ret elif isinstance(data, tuple): - tuple_ret = tuple(convert_to_tensor(i, dtype=dtype, device=device) for i in data) - return torch.as_tensor(tuple_ret, dtype=dtype, device=device) if wrap_sequence else tuple_ret # type: ignore + tuple_ret = tuple(convert_to_tensor(i, dtype=dtype, device=device, track_meta=track_meta) for i in data) + return _convert_tensor(tuple_ret, dtype=dtype, device=device) if wrap_sequence else tuple_ret elif isinstance(data, dict): - return {k: convert_to_tensor(v, dtype=dtype, device=device) for k, v in data.items()} + return {k: convert_to_tensor(v, dtype=dtype, device=device, track_meta=track_meta) for k, v in data.items()} return data @@ -155,7 +172,7 @@ def convert_to_numpy(data, dtype: DtypeLike = None, wrap_sequence: bool = False) E.g., `[1, 2]` -> `[array(1), array(2)]`. If `True`, then `[1, 2]` -> `array([1, 2])`. """ if isinstance(data, torch.Tensor): - data = data.detach().to(dtype=get_equivalent_dtype(dtype, torch.Tensor), device="cpu").numpy() + data = np.asarray(data.detach().to(device="cpu").numpy(), dtype=get_equivalent_dtype(dtype, np.ndarray)) elif has_cp and isinstance(data, cp_ndarray): data = cp.asnumpy(data).astype(dtype, copy=False) elif isinstance(data, (np.ndarray, float, int, bool)): @@ -218,17 +235,19 @@ def convert_data_type( wrap_sequence: bool = False, ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]: """ - Convert to `torch.Tensor`/`np.ndarray` from `torch.Tensor`/`np.ndarray`/`float`/`int` etc. + Convert to `MetaTensor`, `torch.Tensor` or `np.ndarray` from `MetaTensor`, `torch.Tensor`, + `np.ndarray`, `float`, `int`, etc. Args: data: data to be converted - output_type: `torch.Tensor` or `np.ndarray` (if `None`, unchanged) - device: if output is `torch.Tensor`, select device (if `None`, unchanged) + output_type: `monai.data.MetaTensor`, `torch.Tensor`, or `np.ndarray` (if `None`, unchanged) + device: if output is `MetaTensor` or `torch.Tensor`, select device (if `None`, unchanged) dtype: dtype of output data. Converted to correct library type (e.g., `np.float32` is converted to `torch.float32` if output type is `torch.Tensor`). If left blank, it remains unchanged. wrap_sequence: if `False`, then lists will recursively call this function. E.g., `[1, 2]` -> `[array(1), array(2)]`. If `True`, then `[1, 2]` -> `array([1, 2])`. + Returns: modified data, orig_type, orig_device @@ -242,7 +261,9 @@ def convert_data_type( """ orig_type: type - if isinstance(data, torch.Tensor): + if isinstance(data, monai.data.MetaTensor): + orig_type = monai.data.MetaTensor + elif isinstance(data, torch.Tensor): orig_type = torch.Tensor elif isinstance(data, np.ndarray): orig_type = np.ndarray @@ -258,8 +279,10 @@ def convert_data_type( dtype_ = get_equivalent_dtype(dtype, output_type) data_: NdarrayTensor + if issubclass(output_type, torch.Tensor): - data_ = convert_to_tensor(data, dtype=dtype_, device=device, wrap_sequence=wrap_sequence) + track_meta = issubclass(output_type, monai.data.MetaTensor) + data_ = convert_to_tensor(data, dtype=dtype_, device=device, wrap_sequence=wrap_sequence, track_meta=track_meta) return data_, orig_type, orig_device if issubclass(output_type, np.ndarray): data_ = convert_to_numpy(data, dtype=dtype_, wrap_sequence=wrap_sequence) @@ -289,15 +312,39 @@ def convert_to_dst_type( See Also: :func:`convert_data_type` """ + device = dst.device if isinstance(dst, torch.Tensor) else None if dtype is None: dtype = dst.dtype + copy_meta = False output_type: Any - if isinstance(dst, torch.Tensor): + if isinstance(dst, monai.data.MetaTensor): + output_type = monai.data.MetaTensor + if not isinstance(src, monai.data.MetaTensor): + copy_meta = True # converting a non-meta tensor to a meta tensor, probably take the metadata as well. + elif isinstance(dst, torch.Tensor): output_type = torch.Tensor elif isinstance(dst, np.ndarray): output_type = np.ndarray else: output_type = type(dst) - return convert_data_type(data=src, output_type=output_type, device=device, dtype=dtype, wrap_sequence=wrap_sequence) + output: NdarrayTensor + output, _type, _device = convert_data_type( + data=src, output_type=output_type, device=device, dtype=dtype, wrap_sequence=wrap_sequence + ) + if copy_meta and isinstance(output, monai.data.MetaTensor): + output.meta, output.applied_operations = deepcopy(dst.meta), deepcopy(dst.applied_operations) # type: ignore + return output, _type, _device + + +def convert_to_list(data: Union[Sequence, torch.Tensor, np.ndarray]) -> list: + """ + Convert to list from `torch.Tensor`/`np.ndarray`/`list`/`tuple` etc. + Args: + data: data to be converted + Returns: + a list + + """ + return data.tolist() if isinstance(data, (torch.Tensor, np.ndarray)) else list(data) diff --git a/monai/visualize/__init__.py b/monai/visualize/__init__.py index cd980846b3..49a628b66f 100644 --- a/monai/visualize/__init__.py +++ b/monai/visualize/__init__.py @@ -10,6 +10,7 @@ # limitations under the License. from .class_activation_maps import CAM, GradCAM, GradCAMpp, ModelWithHooks, default_normalizer +from .gradient_based import GuidedBackpropGrad, GuidedBackpropSmoothGrad, SmoothGrad, VanillaGrad from .img2tensorboard import add_animated_gif, make_animated_gif_summary, plot_2d_or_3d_image from .occlusion_sensitivity import OcclusionSensitivity from .utils import blend_images, matshow3d diff --git a/monai/visualize/class_activation_maps.py b/monai/visualize/class_activation_maps.py index 16fb64cb46..ba1f5d2589 100644 --- a/monai/visualize/class_activation_maps.py +++ b/monai/visualize/class_activation_maps.py @@ -89,7 +89,7 @@ def __init__( mod.register_backward_hook(self.backward_hook(name)) if self.register_forward: mod.register_forward_hook(self.forward_hook(name)) - if len(_registered) != len(self.target_layers): + if self.target_layers and (len(_registered) != len(self.target_layers)): warnings.warn(f"Not all target_layers exist in the network module: targets: {self.target_layers}.") def backward_hook(self, name): @@ -139,10 +139,10 @@ def __call__(self, x, class_idx=None, retain_graph=False): self.score.sum().backward(retain_graph=retain_graph) for layer in self.target_layers: if layer not in self.gradients: - raise RuntimeError( + warnings.warn( f"Backward hook for {layer} is not triggered; `requires_grad` of {layer} should be `True`." ) - grad = tuple(self.gradients[layer] for layer in self.target_layers) + grad = tuple(self.gradients[layer] for layer in self.target_layers if layer in self.gradients) if train: self.model.train() return logits, acti, grad diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py new file mode 100644 index 0000000000..32b8110b6d --- /dev/null +++ b/monai/visualize/gradient_based.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 + +from functools import partial +from typing import Callable + +import torch + +from monai.networks.utils import replace_modules_temp +from monai.utils.module import optional_import +from monai.visualize.class_activation_maps import ModelWithHooks + +trange, has_trange = optional_import("tqdm", name="trange") + + +__all__ = ["VanillaGrad", "SmoothGrad", "GuidedBackpropGrad", "GuidedBackpropSmoothGrad"] + + +class _AutoGradReLU(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + pos_mask = (x > 0).type_as(x) + output = torch.mul(x, pos_mask) + ctx.save_for_backward(x, output) + return output + + @staticmethod + def backward(ctx, grad_output): + x, _ = ctx.saved_tensors + pos_mask_1 = (x > 0).type_as(grad_output) + pos_mask_2 = (grad_output > 0).type_as(grad_output) + y = torch.mul(grad_output, pos_mask_1) + grad_input = torch.mul(y, pos_mask_2) + return grad_input + + +class _GradReLU(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + out: torch.Tensor = _AutoGradReLU.apply(x) + return out + + +class VanillaGrad: + def __init__(self, model: torch.nn.Module) -> None: + if not isinstance(model, ModelWithHooks): # Convert to model with hooks if necessary + self._model = ModelWithHooks(model, target_layer_names=(), register_backward=True) + else: + self._model = model + + @property + def model(self): + return self._model.model + + @model.setter + def model(self, m): + if not isinstance(m, ModelWithHooks): # regular model as ModelWithHooks + self._model.model = m + else: + self._model = m # replace the ModelWithHooks + + def get_grad(self, x: torch.Tensor, index: torch.Tensor | int | None, retain_graph=True) -> torch.Tensor: + if x.shape[0] != 1: + raise ValueError("expect batch size of 1") + x.requires_grad = True + + self._model(x, class_idx=index, retain_graph=retain_graph) + grad: torch.Tensor = x.grad.detach() + return grad + + def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None) -> torch.Tensor: + return self.get_grad(x, index) + + +class SmoothGrad(VanillaGrad): + """ + See also: + - Smilkov et al. SmoothGrad: removing noise by adding noise https://arxiv.org/abs/1706.03825 + """ + + def __init__( + self, + model: torch.nn.Module, + stdev_spread: float = 0.15, + n_samples: int = 25, + magnitude: bool = True, + verbose: bool = True, + ) -> None: + super().__init__(model) + self.stdev_spread = stdev_spread + self.n_samples = n_samples + self.magnitude = magnitude + self.range: Callable + if verbose and has_trange: + self.range = partial(trange, desc=f"Computing {self.__class__.__name__}") + else: + self.range = range + + def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None) -> torch.Tensor: + stdev = (self.stdev_spread * (x.max() - x.min())).item() + total_gradients = torch.zeros_like(x) + for _ in self.range(self.n_samples): + # create noisy image + noise = torch.normal(0, stdev, size=x.shape, dtype=torch.float32, device=x.device) + x_plus_noise = x + noise + x_plus_noise = x_plus_noise.detach() + + # get gradient and accumulate + grad = self.get_grad(x_plus_noise, index) + total_gradients += (grad * grad) if self.magnitude else grad + + # average + if self.magnitude: + total_gradients = total_gradients**0.5 + + return total_gradients / self.n_samples + + +class GuidedBackpropGrad(VanillaGrad): + def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None) -> torch.Tensor: + with replace_modules_temp(self.model, "relu", _GradReLU(), strict_match=False): + return super().__call__(x, index) + + +class GuidedBackpropSmoothGrad(SmoothGrad): + def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None) -> torch.Tensor: + with replace_modules_temp(self.model, "relu", _GradReLU(), strict_match=False): + return super().__call__(x, index) diff --git a/monai/visualize/img2tensorboard.py b/monai/visualize/img2tensorboard.py index 0af05adf32..6dd7abcbf1 100644 --- a/monai/visualize/img2tensorboard.py +++ b/monai/visualize/img2tensorboard.py @@ -21,14 +21,13 @@ PIL, _ = optional_import("PIL") GifImage, _ = optional_import("PIL.GifImagePlugin", name="Image") SummaryX, _ = optional_import("tensorboardX.proto.summary_pb2", name="Summary") +SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") SummaryWriterX, has_tensorboardx = optional_import("tensorboardX", name="SummaryWriter") if TYPE_CHECKING: from tensorboard.compat.proto.summary_pb2 import Summary - from torch.utils.tensorboard import SummaryWriter else: Summary, _ = optional_import("tensorboard.compat.proto.summary_pb2", name="Summary") - SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") __all__ = ["make_animated_gif_summary", "add_animated_gif", "plot_2d_or_3d_image"] @@ -104,7 +103,7 @@ def make_animated_gif_summary( def add_animated_gif( - writer: SummaryWriter, + writer, tag: str, image_tensor: Union[np.ndarray, torch.Tensor], max_out: int = 3, @@ -136,7 +135,7 @@ def add_animated_gif( def plot_2d_or_3d_image( data: Union[NdarrayTensor, List[NdarrayTensor]], step: int, - writer: SummaryWriter, + writer, index: int = 0, max_channels: int = 1, frame_dim: int = -3, diff --git a/monai/visualize/utils.py b/monai/visualize/utils.py index c1111abd82..1ef6d6da57 100644 --- a/monai/visualize/utils.py +++ b/monai/visualize/utils.py @@ -52,7 +52,7 @@ def matshow3d( Higher dimensional arrays will be reshaped into (-1, H, W, [C]), `C` depends on `channel_dim` arg. A list of channel-first (C, H[, W, D]) arrays can also be passed in, in which case they will be displayed as a padded and stacked volume. - fig: matplotlib figure to use. If None, a new figure will be created. + fig: matplotlib figure or Axes to use. If None, a new figure will be created. title: title of the figure. figsize: size of the figure. frames_per_row: number of frames to display in each row. If None, sqrt(firstdim) will be used. @@ -136,17 +136,20 @@ def matshow3d( im = np.moveaxis(im, 0, -1) # figure related configurations - if fig is None: - fig = plt.figure(tight_layout=True) - if not fig.axes: - fig.add_subplot(111) - ax = fig.axes[0] + if isinstance(fig, plt.Axes): + ax = fig + else: + if fig is None: + fig = plt.figure(tight_layout=True) + if not fig.axes: + fig.add_subplot(111) + ax = fig.axes[0] ax.matshow(im, vmin=vmin, vmax=vmax, interpolation=interpolation, **kwargs) ax.axis("off") if title is not None: ax.set_title(title) - if figsize is not None: + if figsize is not None and hasattr(fig, "set_size_inches"): fig.set_size_inches(figsize) if show: plt.show() diff --git a/pyproject.toml b/pyproject.toml index 03e9f49ab5..59db27e134 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ requires = [ "wheel", "setuptools", - "torch>=1.6", + "torch>=1.7", "ninja", ] @@ -30,3 +30,7 @@ exclude = ''' | monai/_version.py ) ''' + +[tool.pycln] +all = true +exclude = "monai/bundle/__main__.py" diff --git a/requirements-dev.txt b/requirements-dev.txt index 651a99eba9..8318a1795c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,13 +1,13 @@ # Full requirements for developments -r requirements-min.txt -pytorch-ignite==0.4.8 -gdown>=3.6.4 +pytorch-ignite==0.4.9 +gdown>=4.4.0 scipy itk>=5.2 nibabel pillow!=8.3.0 # https://github.com/python-pillow/Pillow/issues/5571 tensorboard -scikit-image>=0.14.2 +scikit-image>=0.19.0 tqdm>=4.47.0 lmdb flake8>=3.8.1 @@ -47,3 +47,6 @@ types-PyYAML pyyaml fire jsonschema +pynrrd +pre-commit +pydicom diff --git a/requirements.txt b/requirements.txt index e4ea34b5d4..14eb2b30e9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -torch>=1.6 +torch>=1.7 numpy>=1.17 diff --git a/runtests.sh b/runtests.sh index 5464f3d020..9e6ef3d0e1 100755 --- a/runtests.sh +++ b/runtests.sh @@ -14,6 +14,13 @@ # script for running all tests set -e +# FIXME: https://github.com/Project-MONAI/MONAI/issues/4354 +protobuf_major_version=$(pip list | grep '^protobuf ' | tr -s ' ' | cut -d' ' -f2 | cut -d'.' -f1) +if [ "$protobuf_major_version" -ge "4" ] +then + export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python +fi + # output formatting separator="" blue="" @@ -51,6 +58,7 @@ doPytypeFormat=false doMypyFormat=false doCleanup=false doDistTests=false +doPrecommit=false NUM_PARALLEL=1 @@ -59,7 +67,7 @@ PY_EXE=${MONAI_PY_EXE:-$(which python)} function print_usage { echo "runtests.sh [--codeformat] [--autofix] [--black] [--isort] [--flake8] [--pylint] [--clangformat] [--pytype] [--mypy]" echo " [--unittests] [--disttests] [--coverage] [--quick] [--min] [--net] [--dryrun] [-j number] [--list_tests]" - echo " [--copyright] [--build] [--clean] [--help] [--version]" + echo " [--copyright] [--build] [--clean] [--precommit] [--help] [--version]" echo "" echo "MONAI unit testing utilities." echo "" @@ -78,6 +86,7 @@ function print_usage { echo " --flake8 : perform \"flake8\" code format checks" echo " --pylint : perform \"pylint\" code format checks" echo " --clangformat : format csrc code using \"clang-format\"" + 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" @@ -146,6 +155,8 @@ function clang_format { exit 1 fi find monai/csrc -type f | while read i; do $clang_format_tool -style=file -i $i; done + find monai/_extensions -type f -name "*.cpp" -o -name "*.h" -o -name "*.cuh" -o -name "*.cu" |\ + while read i; do $clang_format_tool -style=file -i $i; done } function clean_py { @@ -273,6 +284,9 @@ do --pylint) doPylintFormat=true ;; + --precommit) + doPrecommit=true + ;; --pytype) doPytypeFormat=true ;; @@ -359,7 +373,6 @@ then clang_format echo "${green}done!${noColor}" - exit fi # unconditionally report on the state of monai @@ -390,6 +403,30 @@ then fi +if [ $doPrecommit = true ] +then + set +e # disable exit on failure so that diagnostics can be given on failure + echo "${separator}${blue}pre-commit${noColor}" + + # ensure that the necessary packages for code format testing are installed + if ! is_pip_installed pre_commit + then + install_deps + fi + ${cmdPrefix}${PY_EXE} -m pre_commit run --all-files + + pre_commit_status=$? + if [ ${pre_commit_status} -ne 0 ] + then + print_style_fail_msg + exit ${pre_commit_status} + else + echo "${green}passed!${noColor}" + fi + set -e # enable exit on failure +fi + + if [ $doIsortFormat = true ] then set +e # disable exit on failure so that diagnostics can be given on failure @@ -405,7 +442,7 @@ then then install_deps fi - ${cmdPrefix}isort --version + ${cmdPrefix}${PY_EXE} -m isort --version if [ $doIsortFix = true ] then @@ -499,8 +536,8 @@ then fi ${cmdPrefix}${PY_EXE} -m pylint --version - ignore_codes="E1101,E1102,E0601,E1130,E1123,E0102,E1120,E1137,E1136" - ${cmdPrefix}${PY_EXE} -m pylint monai tests -E --disable=$ignore_codes -j $NUM_PARALLEL + ignore_codes="C,R,W,E1101,E1102,E0601,E1130,E1123,E0102,E1120,E1137,E1136" + ${cmdPrefix}${PY_EXE} -m pylint monai tests --disable=$ignore_codes -j $NUM_PARALLEL pylint_status=$? if [ ${pylint_status} -ne 0 ] diff --git a/setup.cfg b/setup.cfg index a7d597d6bd..9a5efba82d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -24,7 +24,7 @@ setup_requires = torch ninja install_requires = - torch>=1.6 + torch>=1.7 numpy>=1.17 [options.extras_require] @@ -33,8 +33,8 @@ all = scikit-image>=0.14.2 pillow tensorboard - gdown>=3.6.4 - pytorch-ignite==0.4.8 + gdown>=4.4.0 + pytorch-ignite==0.4.9 torchvision itk>=5.2 tqdm>=4.47.0 @@ -53,6 +53,8 @@ all = pyyaml fire jsonschema + pynrrd + pydicom nibabel = nibabel skimage = @@ -62,9 +64,9 @@ pillow = tensorboard = tensorboard gdown = - gdown>=3.6.4 + gdown>=4.4.0 ignite = - pytorch-ignite==0.4.8 + pytorch-ignite==0.4.9 torchvision = torchvision itk = @@ -101,6 +103,10 @@ fire = fire jsonschema = jsonschema +pynrrd = + pynrrd +pydicom = + pydicom [flake8] select = B,C,E,F,N,P,T4,W,B9 @@ -115,6 +121,7 @@ ignore = W504 C408 N812 # lowercase 'torch.nn.functional' imported as non lowercase 'F' + B023 # https://github.com/Project-MONAI/MONAI/issues/4627 per_file_ignores = __init__.py: F401, __main__.py: F401 exclude = *.pyi,.git,.eggs,monai/_version.py,versioneer.py,venv,.venv,_version.py diff --git a/tests/croppers.py b/tests/croppers.py new file mode 100644 index 0000000000..8f78249d90 --- /dev/null +++ b/tests/croppers.py @@ -0,0 +1,103 @@ +# 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. + +import unittest +from copy import deepcopy + +import numpy as np + +from monai.data.meta_tensor import MetaTensor +from monai.transforms.transform import MapTransform +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose + + +class CropTest(unittest.TestCase): + @staticmethod + def get_arr(shape): + return np.random.randint(100, size=shape).astype(float) + + def crop_test(self, input_param, input_shape, expected_shape, same_area=None): + base_comparison = None + input_image = self.get_arr(input_shape) + + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + # input parameters, such as roi_start can be numpy, torch, list etc. + for param_type in TEST_NDARRAYS_ALL + (None,): + with self.subTest(param_type=param_type): + input_param_mod = deepcopy(input_param) + if param_type is not None: + for k in ("roi_start", "roi_end", "roi_center", "roi_size", "roi_scale"): + if k in input_param: + input_param_mod[k] = param_type(input_param[k]) + im = im_type(input_image) + cropper = self.Cropper(**input_param_mod) + is_map = isinstance(cropper, MapTransform) + input_data = {"img": im} if is_map else im + result = cropper(input_data) + out_im = result["img"] if is_map else result + self.assertIsInstance(out_im, MetaTensor) + self.assertTupleEqual(out_im.shape, expected_shape) + if same_area is not None: + assert_allclose(out_im, im[same_area], type_test=False) + # check result is the same regardless of input type + if base_comparison is None: + base_comparison = out_im + else: + assert_allclose(out_im, base_comparison) + + # test inverse + inv = cropper.inverse(result) + inv_im = inv["img"] if is_map else inv + self.assertIsInstance(inv_im, MetaTensor) + if same_area is not None: + assert_allclose(inv_im[same_area], im[same_area], type_test=False) + self.assertEqual(inv_im.applied_operations, []) + + def crop_test_value(self, input_param, input_arr, expected_array): + cropper = self.Cropper(**input_param) + is_map = isinstance(cropper, MapTransform) + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + im = im_type(input_arr) + input_data = {"img": im} if is_map else im + result = self.Cropper(**input_param)(input_data) + out_im = result["img"] if is_map else result + self.assertIsInstance(out_im, MetaTensor) + assert_allclose(out_im, expected_array, type_test=False) + + def multi_inverse(self, input_shape, init_params): + input_data = np.arange(np.prod(input_shape)).reshape(*input_shape) + 1 + xform = self.Cropper(**init_params) + xform.set_random_state(1234) + out = xform(input_data) + if "num_samples" in init_params: + self.assertEqual(len(out), init_params["num_samples"]) + inv = xform.inverse(out) + self.assertIsInstance(inv, MetaTensor) + self.assertEqual(inv.applied_operations, []) + self.assertTrue("patch_index" not in inv.meta) + self.assertTupleEqual(inv.shape, input_shape) + inv_np = inv.numpy() + + # get list of all numbers that exist inside the crops + uniques = set() + for o in out: + uniques.update(set(o.flatten().tolist())) + + # make sure that + for i in uniques: + a = np.where(input_data == i) + b = np.where(inv_np == i) + self.assertTupleEqual(a, b) + # there should be as many zeros as elements missing from uniques + missing = input_data.size - len(uniques) + self.assertEqual((inv_np == 0).sum(), missing) diff --git a/tests/min_tests.py b/tests/min_tests.py index 9bf95f3f49..898d1b7b00 100644 --- a/tests/min_tests.py +++ b/tests/min_tests.py @@ -41,6 +41,7 @@ def run_testsuit(): "test_dataset", "test_dataset_summary", "test_deepedit_transforms", + "test_deepedit_interaction", "test_deepgrow_dataset", "test_deepgrow_interaction", "test_deepgrow_transforms", @@ -52,6 +53,8 @@ def run_testsuit(): "test_ensure_channel_firstd", "test_fill_holes", "test_fill_holesd", + "test_foreground_mask", + "test_foreground_maskd", "test_global_mutual_information_loss", "test_handler_checkpoint_loader", "test_handler_checkpoint_saver", @@ -76,7 +79,6 @@ def run_testsuit(): "test_handler_regression_metrics_dist", "test_handler_rocauc", "test_handler_rocauc_dist", - "test_handler_segmentation_saver", "test_handler_smartcache", "test_handler_stats", "test_handler_surface_distance", @@ -113,6 +115,8 @@ def run_testsuit(): "test_nifti_header_revise", "test_nifti_rw", "test_nifti_saver", + "test_nuclick_transforms", + "test_nrrd_reader", "test_occlusion_sensitivity", "test_orientation", "test_orientationd", @@ -143,7 +147,9 @@ def run_testsuit(): "test_smartcachedataset", "test_spacing", "test_spacingd", + "test_splitdimd", "test_surface_distance", + "test_surface_dice", "test_testtimeaugmentation", "test_torchvision", "test_torchvisiond", @@ -155,6 +161,7 @@ def run_testsuit(): "test_vitautoenc", "test_write_metrics_reports", "test_wsireader", + "test_wsireader_new", "test_zoom", "test_zoom_affine", "test_zoomd", @@ -163,6 +170,8 @@ def run_testsuit(): "test_bundle_verify_metadata", "test_bundle_verify_net", "test_bundle_ckpt_export", + "test_bundle_utils", + "test_bundle_init_bundle", ] assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}" diff --git a/tests/ngc_mmar_loading.py b/tests/ngc_mmar_loading.py index 5261dfd612..4bb89dde0e 100644 --- a/tests/ngc_mmar_loading.py +++ b/tests/ngc_mmar_loading.py @@ -10,6 +10,7 @@ # limitations under the License. import os +import sys import unittest import torch @@ -17,6 +18,7 @@ from monai.apps.mmars import MODEL_DESC, load_from_mmar from monai.config import print_debug_info +from monai.networks.utils import copy_model_state class TestAllDownloadingMMAR(unittest.TestCase): @@ -26,10 +28,33 @@ def setUp(self): @parameterized.expand((item,) for item in MODEL_DESC) def test_loading_mmar(self, item): + if item["name"] == "clara_pt_self_supervised_learning_segmentation": # test the byow model + default_model_file = os.path.join("ssl_models_2gpu", "best_metric_model.pt") + pretrained_weights = load_from_mmar( + item=item["name"], + mmar_dir="./", + map_location="cpu", + api=True, + model_file=default_model_file, + weights_only=True, + ) + pretrained_weights = {k.split(".", 1)[1]: v for k, v in pretrained_weights["state_dict"].items()} + sys.path.append(os.path.join(f"{item['name']}", "custom")) # custom model folder + from vit_network import ViTAutoEnc # pylint: disable=E0401 + + model = ViTAutoEnc( + in_channels=1, + img_size=(96, 96, 96), + patch_size=(16, 16, 16), + pos_embed="conv", + hidden_size=768, + mlp_dim=3072, + ) + _, loaded, not_loaded = copy_model_state(model, pretrained_weights) + self.assertTrue(len(loaded) > 0 and len(not_loaded) == 0) + return if item["name"] == "clara_pt_fed_learning_brain_tumor_mri_segmentation": default_model_file = os.path.join("models", "server", "best_FL_global_model.pt") - elif item["name"] == "clara_pt_self_supervised_learning_segmentation": - default_model_file = os.path.join("models_2gpu", "best_metric_model.pt") else: default_model_file = None pretrained_model = load_from_mmar( diff --git a/tests/padders.py b/tests/padders.py new file mode 100644 index 0000000000..3fa3280cb5 --- /dev/null +++ b/tests/padders.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. + +import unittest +from typing import List + +import numpy as np +import torch + +from monai.data.meta_tensor import MetaTensor +from monai.transforms.transform import MapTransform +from monai.utils.enums import NumpyPadMode, PytorchPadMode +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose + +MODES = [] +# Test modes +NP_MODES: List = [ + "constant", + "edge", + # `reflect` mode is not supported in some PyTorch versions, skip the test + # "reflect", + "wrap", + "median", + "mean", +] +MODES += NP_MODES +MODES += [NumpyPadMode(i) for i in NP_MODES] + +PT_MODES: list = [ + "constant", + "replicate", + "circular", + # `reflect` mode is not supported in some PyTorch versions, skip the test + # "reflect", +] +MODES += PT_MODES +MODES += [PytorchPadMode(i) for i in PT_MODES] + + +class PadTest(unittest.TestCase): + @staticmethod + def get_arr(shape): + return np.random.randint(100, size=shape).astype(float) + + def pad_test(self, input_param, input_shape, expected_shape, modes=None): + # loop over each mode + for mode in modes or MODES: + with self.subTest(mode=mode): + base_comparison = None + im = self.get_arr(input_shape) + padder = self.Padder(mode=mode, **input_param) + is_map = isinstance(padder, MapTransform) + # check result is the same regardless of input type + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + input_image = im_type(im) + input_data = {"img": im_type(im)} if is_map else im_type(im) + # our array transforms can also take `mode` as an argument to `__call__` + # Check this gives equivalent results + for call_extra_args in [{}] if is_map else [{}, {"mode": mode}]: + with self.subTest(call_extra_args=call_extra_args): + r_out = padder(input_data, **call_extra_args) + r_im = r_out["img"] if is_map else r_out + # check shape, type, etc. + np.testing.assert_allclose(r_im.shape, expected_shape) + self.assertIsInstance(r_im, MetaTensor) + self.assertEqual(len(r_im.applied_operations), 1) + # check results are same regardless of input type + if base_comparison is None: + base_comparison = r_im + else: + assert_allclose(r_im, base_comparison) + # test inverse + if isinstance(r_im, MetaTensor): + r_out = padder.inverse(r_out) + r_im = r_out["img"] if is_map else r_out + self.assertIsInstance(r_im, MetaTensor) + assert_allclose(r_im, input_image, type_test=False) + self.assertEqual(r_im.applied_operations, []) + + def pad_test_kwargs(self, unchanged_slices, **input_param): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + for kwargs in ({"value": 2}, {"constant_values": ((0, 0), (1, 1), (2, 2))}): + with self.subTest(kwargs=kwargs): + im = im_type(np.random.randint(-100, -10, size=(3, 8, 4))) + padder = self.Padder(**input_param, **kwargs) + result = padder(im) + if isinstance(result, torch.Tensor): + result = result.cpu() + assert_allclose(result[unchanged_slices], im, type_test=False) + # we should have the same as the input plus some 2s (if value) or 1s and 2s (if constant_values) + expected_vals = np.unique(im).tolist() + expected_vals += [2] if "value" in kwargs else [1, 2] + assert_allclose(np.unique(result), expected_vals, type_test=False) + # check inverse + if isinstance(result, MetaTensor): + inv = padder.inverse(result) + assert_allclose(im, inv, type_test=False) + self.assertEqual(inv.applied_operations, []) diff --git a/tests/profile_subclass/README.md b/tests/profile_subclass/README.md new file mode 100644 index 0000000000..de16ef2d91 --- /dev/null +++ b/tests/profile_subclass/README.md @@ -0,0 +1,43 @@ +# Profiling the performance of subclassing/`__torch_function__` in MONAI + +## Requirements +```bash +pip install py-spy +pip install snakeviz # for viewing the cProfile results +``` + +## Commands + +### Install MONAI +``` +./runtests.sh --build # from monai's root directory +``` +or follow the installation guide (https://docs.monai.io/en/latest/installation.html) + +### Profiling the task of adding two MetaTensors +```bash +python profiling.py +``` + +### Profiling using `py-spy` +```bash +py-spy record -o Tensor.svg -- python pyspy_profiling.py Tensor +py-spy record -o SubTensor.svg -- python pyspy_profiling.py SubTensor +py-spy record -o SubWithTorchFunc.svg -- python pyspy_profiling.py SubWithTorchFunc +py-spy record -o MetaTensor.svg -- python pyspy_profiling.py MetaTensor +``` + +### Profiling using `cProfile` and `SNAKEVIZ` + +```bash +python cprofile_profiling.py +snakeviz out_200.prof +``` + +--- +These tests are based on the following code: +https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark + +- Overhead for torch functions when run on `torch.Tensor` objects is on the order of 2 microseconds. +- `__torch_function__` should add zero overhead for `torch.Tensor` inputs, a small overhead for subclasses of `torch.Tensor`, and an order of microseconds for `MeatTensor`. +- Changing the dispatching mechanism may result in changes that are on the order of 100 ns, which are hard to detect due to noise, but important. diff --git a/tests/profile_subclass/cprofile_profiling.py b/tests/profile_subclass/cprofile_profiling.py new file mode 100644 index 0000000000..a6c940c9c0 --- /dev/null +++ b/tests/profile_subclass/cprofile_profiling.py @@ -0,0 +1,28 @@ +# 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. + +""" +Profiling MetaTensor +""" + +import cProfile + +import torch + +from monai.data.meta_tensor import MetaTensor + +if __name__ == "__main__": + n_chan = 3 + for hwd in (10, 200): + shape = (n_chan, hwd, hwd, hwd) + a = MetaTensor(torch.rand(shape), meta={"affine": torch.eye(4) * 2, "fname": "something1"}) + b = MetaTensor(torch.rand(shape), meta={"affine": torch.eye(4) * 3, "fname": "something2"}) + cProfile.run("c = a + b", filename=f"out_{hwd}.prof") diff --git a/tests/profile_subclass/min_classes.py b/tests/profile_subclass/min_classes.py new file mode 100644 index 0000000000..87c0ce671d --- /dev/null +++ b/tests/profile_subclass/min_classes.py @@ -0,0 +1,29 @@ +# 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. + + +""" +Minimal subclassing as baselines +Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark +""" + +import torch + +__all__ = ["SubTensor", "SubWithTorchFunc"] + + +class SubTensor(torch.Tensor): + pass + + +class SubWithTorchFunc(torch.Tensor): + def __torch_function__(self, func, types, args=(), kwargs=None): + return super().__torch_function__(func, types, args, {} if kwargs is None else kwargs) diff --git a/tests/profile_subclass/profiling.py b/tests/profile_subclass/profiling.py new file mode 100644 index 0000000000..28740e82e1 --- /dev/null +++ b/tests/profile_subclass/profiling.py @@ -0,0 +1,72 @@ +# 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. + +""" +Comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor +Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark +""" +import argparse + +import torch +from min_classes import SubTensor, SubWithTorchFunc + +from monai.data import MetaTensor +from monai.utils.profiling import PerfContext + +NUM_REPEATS = 1000 +NUM_REPEAT_OF_REPEATS = 1000 + + +def bench(t1, t2): + bench_times = [] + for _ in range(NUM_REPEAT_OF_REPEATS): + with PerfContext() as pc: + for _ in range(NUM_REPEATS): + torch.add(t1, t2) + bench_times.append(pc.total_time) + + bench_time_min = float(torch.min(torch.Tensor(bench_times))) / NUM_REPEATS + bench_time_avg = float(torch.sum(torch.Tensor(bench_times))) / (NUM_REPEATS * NUM_REPEAT_OF_REPEATS) + bench_time_med = float(torch.median(torch.Tensor(bench_times))) / NUM_REPEATS + bench_std = float(torch.std(torch.Tensor(bench_times))) / NUM_REPEATS + return bench_time_min, bench_time_avg, bench_time_med, bench_std + + +def main(): + global NUM_REPEATS + global NUM_REPEAT_OF_REPEATS + + parser = argparse.ArgumentParser(description="Run the __torch_function__ benchmarks.") + parser.add_argument( + "--nreps", "-n", type=int, default=NUM_REPEATS, help="The number of repeats for one measurement." + ) + parser.add_argument("--nrepreps", "-m", type=int, default=NUM_REPEAT_OF_REPEATS, help="The number of measurements.") + args = parser.parse_args() + + NUM_REPEATS = args.nreps + NUM_REPEAT_OF_REPEATS = args.nrepreps + + types = torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor + + for t in types: + tensor_1 = t(1) + tensor_2 = t(2) + + b_min, b_avg, b_med, b_std = bench(tensor_1, tensor_2) + print( + "Type {} time (microseconds): min: {}, avg: {}, median: {}, and std {}.".format( + t.__name__, (10**6 * b_min), (10**6) * b_avg, (10**6) * b_med, (10**6) * b_std + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/profile_subclass/pyspy_profiling.py b/tests/profile_subclass/pyspy_profiling.py new file mode 100644 index 0000000000..302bfd39c3 --- /dev/null +++ b/tests/profile_subclass/pyspy_profiling.py @@ -0,0 +1,40 @@ +# 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. + +""" +To be used with py-spy, comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor +Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark +""" +import argparse + +import torch +from min_classes import SubTensor, SubWithTorchFunc # noqa: F401 + +from monai.data import MetaTensor # noqa: F401 + +Tensor = torch.Tensor + +NUM_REPEATS = 1000000 + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run the torch.add for a given class a given number of times.") + parser.add_argument("tensor_class", metavar="TensorClass", type=str, help="The class to benchmark.") + parser.add_argument("--nreps", "-n", type=int, default=NUM_REPEATS, help="The number of repeats.") + args = parser.parse_args() + + TensorClass = globals()[args.tensor_class] + NUM_REPEATS = args.nreps + + t1 = TensorClass(1) + t2 = TensorClass(2) + + for _ in range(NUM_REPEATS): + torch.add(t1, t2) diff --git a/tests/test_activations.py b/tests/test_activations.py index a67e6f8cb6..a06316b253 100644 --- a/tests/test_activations.py +++ b/tests/test_activations.py @@ -76,12 +76,12 @@ class TestActivations(unittest.TestCase): - @parameterized.expand(TEST_CASES[:3]) + @parameterized.expand(TEST_CASES) def test_value_shape(self, input_param, img, out, expected_shape): result = Activations(**input_param)(img) def _compare(ret, out, shape): - assert_allclose(ret, out, rtol=1e-3) + assert_allclose(ret, out, rtol=1e-3, type_test=False) self.assertTupleEqual(ret.shape, shape) if isinstance(result, (list, tuple)): diff --git a/tests/test_activationsd.py b/tests/test_activationsd.py index 557d68de90..e38f36e49d 100644 --- a/tests/test_activationsd.py +++ b/tests/test_activationsd.py @@ -51,10 +51,10 @@ class TestActivationsd(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_value_shape(self, input_param, test_input, output, expected_shape): result = Activationsd(**input_param)(test_input) - assert_allclose(result["pred"], output["pred"], rtol=1e-3) + assert_allclose(result["pred"], output["pred"], rtol=1e-3, type_test="tensor") self.assertTupleEqual(result["pred"].shape, expected_shape) if "label" in result: - assert_allclose(result["label"], output["label"], rtol=1e-3) + assert_allclose(result["label"], output["label"], rtol=1e-3, type_test="tensor") self.assertTupleEqual(result["label"].shape, expected_shape) diff --git a/tests/test_add_extreme_points_channel.py b/tests/test_add_extreme_points_channel.py index d2c8a627b6..116f96126f 100644 --- a/tests/test_add_extreme_points_channel.py +++ b/tests/test_add_extreme_points_channel.py @@ -71,7 +71,7 @@ class TestAddExtremePointsChannel(unittest.TestCase): def test_correct_results(self, input_data, expected): add_extreme_points_channel = AddExtremePointsChannel() result = add_extreme_points_channel(**input_data) - assert_allclose(result[IMG_CHANNEL], expected, rtol=1e-4) + assert_allclose(result[IMG_CHANNEL], expected, rtol=1e-4, atol=1e-5) if __name__ == "__main__": diff --git a/tests/test_add_extreme_points_channeld.py b/tests/test_add_extreme_points_channeld.py index 39d221596f..f9837e9ef4 100644 --- a/tests/test_add_extreme_points_channeld.py +++ b/tests/test_add_extreme_points_channeld.py @@ -68,7 +68,7 @@ def test_correct_results(self, input_data, expected): keys="img", label_key="label", sigma=1.0, rescale_min=0.0, rescale_max=1.0 ) result = add_extreme_points_channel(input_data) - assert_allclose(result["img"][IMG_CHANNEL], expected, rtol=1e-4) + assert_allclose(result["img"][IMG_CHANNEL], expected, rtol=1e-4, atol=1e-5) if __name__ == "__main__": diff --git a/tests/test_adjust_contrast.py b/tests/test_adjust_contrast.py index 2f6c4e2259..1c38d0edf3 100644 --- a/tests/test_adjust_contrast.py +++ b/tests/test_adjust_contrast.py @@ -29,7 +29,9 @@ class TestAdjustContrast(NumpyImageTestCase2D): def test_correct_results(self, gamma): adjuster = AdjustContrast(gamma=gamma) for p in TEST_NDARRAYS: - result = adjuster(p(self.imt)) + im = p(self.imt) + result = adjuster(im) + self.assertTrue(type(im), type(result)) if gamma == 1.0: expected = self.imt else: @@ -37,7 +39,7 @@ def test_correct_results(self, gamma): img_min = self.imt.min() img_range = self.imt.max() - img_min expected = np.power(((self.imt - img_min) / float(img_range + epsilon)), gamma) * img_range + img_min - assert_allclose(expected, result, rtol=1e-05, type_test=False) + assert_allclose(result, expected, rtol=1e-05, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_adjust_contrastd.py b/tests/test_adjust_contrastd.py index a7224b643b..2d674c6003 100644 --- a/tests/test_adjust_contrastd.py +++ b/tests/test_adjust_contrastd.py @@ -37,7 +37,7 @@ def test_correct_results(self, gamma): img_min = self.imt.min() img_range = self.imt.max() - img_min expected = np.power(((self.imt - img_min) / float(img_range + epsilon)), gamma) * img_range + img_min - assert_allclose(expected, result["img"], rtol=1e-05, type_test=False) + assert_allclose(result["img"], expected, rtol=1e-05, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_affine.py b/tests/test_affine.py index d681d2941b..019a8f59a4 100644 --- a/tests/test_affine.py +++ b/tests/test_affine.py @@ -10,16 +10,18 @@ # limitations under the License. import unittest +from copy import deepcopy import numpy as np import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Affine -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, test_local_inversion TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -155,11 +157,21 @@ class TestAffine(unittest.TestCase): @parameterized.expand(TESTS) def test_affine(self, input_param, input_data, expected_val): + input_copy = deepcopy(input_data["img"]) g = Affine(**input_param) result = g(**input_data) if isinstance(result, tuple): result = result[0] - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4) + test_local_inversion(g, result, input_copy) + assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4, type_test=False) + + set_track_meta(False) + result = g(**input_data) + if isinstance(result, tuple): + result = result[0] + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_affine_grid.py b/tests/test_affine_grid.py index 6f6364feda..b481601df5 100644 --- a/tests/test_affine_grid.py +++ b/tests/test_affine_grid.py @@ -15,11 +15,12 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import AffineGrid -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, is_tf32_env TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -136,7 +137,11 @@ class TestAffineGrid(unittest.TestCase): @parameterized.expand(TESTS) def test_affine_grid(self, input_param, input_data, expected_val): g = AffineGrid(**input_param) + set_track_meta(False) result, _ = g(**input_data) + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) if "device" in input_data: self.assertEqual(result.device, input_data[device]) assert_allclose(result, expected_val, type_test=False, rtol=_rtol) diff --git a/tests/test_affine_transform.py b/tests/test_affine_transform.py index 5170ab4260..902437350a 100644 --- a/tests/test_affine_transform.py +++ b/tests/test_affine_transform.py @@ -23,12 +23,14 @@ TEST_NORM_CASES = [ [(4, 5), True, [[[0.666667, 0, -1], [0, 0.5, -1], [0, 0, 1]]]], + [(4, 5), True, [[[0.666667, 0, 0], [0, 0.5, 0], [0, 0, 1]]], True], [ (2, 4, 5), True, [[[2.0, 0.0, 0.0, -1.0], [0.0, 0.6666667, 0.0, -1.0], [0.0, 0.0, 0.5, -1.0], [0.0, 0.0, 0.0, 1.0]]], ], [(4, 5), False, [[[0.5, 0.0, -0.75], [0.0, 0.4, -0.8], [0.0, 0.0, 1.0]]]], + [(4, 5), False, [[[0.5, 0.0, 0.25], [0.0, 0.4, 0.2], [0.0, 0.0, 1.0]]], True], [(2, 4, 5), False, [[[1.0, 0.0, 0.0, -0.5], [0.0, 0.5, 0.0, -0.75], [0.0, 0.0, 0.4, -0.8], [0.0, 0.0, 0.0, 1.0]]]], ] @@ -61,6 +63,14 @@ False, [[[1.5, 0.0, 0.0, 0.5], [0.0, 1.25, 0.0, 0.25], [0.0, 0.0, 0.5, -0.5], [0.0, 0.0, 0.0, 1.0]]], ], + [ + [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]], + (2, 4, 6), + (3, 5, 3), + False, + [[[1.5, 0.0, 0.0, 0.0], [0.0, 1.25, 0.0, 0.0], [0.0, 0.0, 0.5, 0.0], [0.0, 0.0, 0.0, 1.0]]], + True, + ], ] TEST_ILL_TO_NORM_AFFINE_CASES = [ @@ -72,15 +82,23 @@ class TestNormTransform(unittest.TestCase): @parameterized.expand(TEST_NORM_CASES) - def test_norm_xform(self, input_shape, align_corners, expected): + def test_norm_xform(self, input_shape, align_corners, expected, zero_centered=False): norm = normalize_transform( - input_shape, device=torch.device("cpu:0"), dtype=torch.float32, align_corners=align_corners + input_shape, + device=torch.device("cpu:0"), + dtype=torch.float32, + align_corners=align_corners, + zero_centered=zero_centered, ) norm = norm.detach().cpu().numpy() np.testing.assert_allclose(norm, expected, atol=1e-6) if torch.cuda.is_available(): norm = normalize_transform( - input_shape, device=torch.device("cuda:0"), dtype=torch.float32, align_corners=align_corners + input_shape, + device=torch.device("cuda:0"), + dtype=torch.float32, + align_corners=align_corners, + zero_centered=zero_centered, ) norm = norm.detach().cpu().numpy() np.testing.assert_allclose(norm, expected, atol=1e-4) @@ -88,15 +106,15 @@ def test_norm_xform(self, input_shape, align_corners, expected): class TestToNormAffine(unittest.TestCase): @parameterized.expand(TEST_TO_NORM_AFFINE_CASES) - def test_to_norm_affine(self, affine, src_size, dst_size, align_corners, expected): + def test_to_norm_affine(self, affine, src_size, dst_size, align_corners, expected, zero_centered=False): affine = torch.as_tensor(affine, device=torch.device("cpu:0"), dtype=torch.float32) - new_affine = to_norm_affine(affine, src_size, dst_size, align_corners) + new_affine = to_norm_affine(affine, src_size, dst_size, align_corners, zero_centered) new_affine = new_affine.detach().cpu().numpy() np.testing.assert_allclose(new_affine, expected, atol=1e-6) if torch.cuda.is_available(): affine = torch.as_tensor(affine, device=torch.device("cuda:0"), dtype=torch.float32) - new_affine = to_norm_affine(affine, src_size, dst_size, align_corners) + new_affine = to_norm_affine(affine, src_size, dst_size, align_corners, zero_centered) new_affine = new_affine.detach().cpu().numpy() np.testing.assert_allclose(new_affine, expected, atol=1e-5, rtol=_rtol) @@ -155,6 +173,13 @@ def test_zoom_2(self): expected = [[[[1, 3]]]] np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol) + def test_zoom_zero_center(self): + affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], dtype=torch.float32) + image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0")) + out = AffineTransform((1, 2), zero_centered=True)(image, affine) + expected = [[[[3, 5]]]] + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol) + def test_affine_transform_minimum(self): t = np.pi / 3 affine = [[np.cos(t), -np.sin(t), 0], [np.sin(t), np.cos(t), 0], [0, 0, 1]] diff --git a/tests/test_affined.py b/tests/test_affined.py index 665c93d23f..b922d80fb5 100644 --- a/tests/test_affined.py +++ b/tests/test_affined.py @@ -10,16 +10,17 @@ # limitations under the License. import unittest +from copy import deepcopy import numpy as np import torch from parameterized import parameterized from monai.transforms import Affined -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, test_local_inversion TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -159,9 +160,11 @@ class TestAffined(unittest.TestCase): @parameterized.expand(TESTS) def test_affine(self, input_param, input_data, expected_val): + input_copy = deepcopy(input_data) g = Affined(**input_param) - result = g(input_data)["img"] - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4) + result = g(input_data) + test_local_inversion(g, result, input_copy, dict_key="img") + assert_allclose(result["img"], expected_val, rtol=1e-4, atol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_anchor_box.py b/tests/test_anchor_box.py new file mode 100644 index 0000000000..a6abbc0200 --- /dev/null +++ b/tests/test_anchor_box.py @@ -0,0 +1,87 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.anchor_utils import AnchorGenerator, AnchorGeneratorWithAnchorShape +from monai.utils import optional_import +from tests.utils import SkipIfBeforePyTorchVersion, assert_allclose, test_script_save + +_, has_torchvision = optional_import("torchvision") + +TEST_CASES_2D = [ + [ + {"sizes": ((10, 12, 14, 16), (20, 24, 28, 32)), "aspect_ratios": ((1.0, 0.5, 2.0), (1.0, 0.5, 2.0))}, + (5, 3, 128, 128), + ((5, 7, 64, 32), (5, 7, 32, 16)), + ] +] + +TEST_CASES_SHAPE_3D = [ + [ + {"feature_map_scales": (1, 2), "base_anchor_shapes": ((4, 3, 6), (8, 2, 4))}, + (5, 3, 128, 128, 128), + ((5, 7, 64, 32, 32), (5, 7, 32, 16, 16)), + ] +] + + +@SkipIfBeforePyTorchVersion((1, 11)) +@unittest.skipUnless(has_torchvision, "Requires torchvision") +class TestAnchorGenerator(unittest.TestCase): + @parameterized.expand(TEST_CASES_2D) + def test_anchor_2d(self, input_param, image_shape, feature_maps_shapes): + + torch_anchor_utils, _ = optional_import("torchvision.models.detection.anchor_utils") + image_list, _ = optional_import("torchvision.models.detection.image_list") + + # test it behaves the same with torchvision for 2d + anchor = AnchorGenerator(**input_param, indexing="xy") + anchor_ref = torch_anchor_utils.AnchorGenerator(**input_param) + for a, a_f in zip(anchor.cell_anchors, anchor_ref.cell_anchors): + assert_allclose(a, a_f, type_test=True, device_test=False, atol=1e-3) + for a, a_f in zip(anchor.num_anchors_per_location(), anchor_ref.num_anchors_per_location()): + assert_allclose(a, a_f, type_test=True, device_test=False, atol=1e-3) + + grid_sizes = [[2, 2], [1, 1]] + strides = [[torch.tensor(1), torch.tensor(2)], [torch.tensor(2), torch.tensor(4)]] + for a, a_f in zip(anchor.grid_anchors(grid_sizes, strides), anchor_ref.grid_anchors(grid_sizes, strides)): + assert_allclose(a, a_f, type_test=True, device_test=False, atol=1e-3) + + images = torch.rand(image_shape) + feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes) + result = anchor(images, feature_maps) + result_ref = anchor_ref(image_list.ImageList(images, ([123, 122],)), feature_maps) + for a, a_f in zip(result, result_ref): + assert_allclose(a, a_f, type_test=True, device_test=False, atol=0.1) + + @parameterized.expand(TEST_CASES_2D) + def test_script_2d(self, input_param, image_shape, feature_maps_shapes): + # test whether support torchscript + anchor = AnchorGenerator(**input_param, indexing="xy") + images = torch.rand(image_shape) + feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes) + test_script_save(anchor, images, feature_maps) + + @parameterized.expand(TEST_CASES_SHAPE_3D) + def test_script_3d(self, input_param, image_shape, feature_maps_shapes): + # test whether support torchscript + anchor = AnchorGeneratorWithAnchorShape(**input_param, indexing="ij") + images = torch.rand(image_shape) + feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes) + test_script_save(anchor, images, feature_maps) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_arraydataset.py b/tests/test_arraydataset.py index ee1a92cf97..380a6372b2 100644 --- a/tests/test_arraydataset.py +++ b/tests/test_arraydataset.py @@ -39,15 +39,17 @@ class TestCompose(Compose): def __call__(self, input_): - img, metadata = self.transforms[0](input_) + img = self.transforms[0](input_) + metadata = img.meta img = self.transforms[1](img) - img, _, _ = self.transforms[2](img, metadata["affine"]) + img = self.transforms[2](img, metadata["affine"]) + metadata = img.meta return self.transforms[3](img), metadata TEST_CASE_3 = [ - TestCompose([LoadImage(image_only=False), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast(prob=1.0)]), - TestCompose([LoadImage(image_only=False), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast(prob=1.0)]), + TestCompose([LoadImage(image_only=True), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast(prob=1.0)]), + TestCompose([LoadImage(image_only=True), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast(prob=1.0)]), (0, 2), (1, 64, 64, 33), ] @@ -58,7 +60,7 @@ def __call__(self, input_): class TestArrayDataset(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) def test_shape(self, img_transform, label_transform, indices, expected_shape): - test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4)) + test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)).astype(float), np.eye(4)) with tempfile.TemporaryDirectory() as tempdir: test_image1 = os.path.join(tempdir, "test_image1.nii.gz") test_seg1 = os.path.join(tempdir, "test_seg1.nii.gz") @@ -92,7 +94,7 @@ def test_shape(self, img_transform, label_transform, indices, expected_shape): @parameterized.expand([TEST_CASE_4]) def test_default_none(self, img_transform, expected_shape): - test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4)) + test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)).astype(float), np.eye(4)) with tempfile.TemporaryDirectory() as tempdir: test_image1 = os.path.join(tempdir, "test_image1.nii.gz") test_image2 = os.path.join(tempdir, "test_image2.nii.gz") @@ -115,7 +117,7 @@ def test_default_none(self, img_transform, expected_shape): @parameterized.expand([TEST_CASE_4]) def test_dataloading_img(self, img_transform, expected_shape): - test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4)) + test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)).astype(float), np.eye(4)) with tempfile.TemporaryDirectory() as tempdir: test_image1 = os.path.join(tempdir, "test_image1.nii.gz") test_image2 = os.path.join(tempdir, "test_image2.nii.gz") @@ -136,7 +138,7 @@ def test_dataloading_img(self, img_transform, expected_shape): @parameterized.expand([TEST_CASE_4]) def test_dataloading_img_label(self, img_transform, expected_shape): - test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4)) + test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)).astype(float), np.eye(4)) with tempfile.TemporaryDirectory() as tempdir: test_image1 = os.path.join(tempdir, "test_image1.nii.gz") test_image2 = os.path.join(tempdir, "test_image2.nii.gz") diff --git a/tests/test_as_channel_first.py b/tests/test_as_channel_first.py index a2d56295b8..732c559a1a 100644 --- a/tests/test_as_channel_first.py +++ b/tests/test_as_channel_first.py @@ -12,10 +12,10 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import AsChannelFirst +from monai.transforms.utils_pytorch_numpy_unification import moveaxis from tests.utils import TEST_NDARRAYS, assert_allclose TESTS = [] @@ -31,10 +31,8 @@ def test_value(self, in_type, input_param, expected_shape): test_data = in_type(np.random.randint(0, 2, size=[1, 2, 3, 4])) result = AsChannelFirst(**input_param)(test_data) self.assertTupleEqual(result.shape, expected_shape) - if isinstance(test_data, torch.Tensor): - test_data = test_data.cpu().numpy() - expected = np.moveaxis(test_data, input_param["channel_dim"], 0) - assert_allclose(result, expected, type_test=False) + expected = moveaxis(test_data, input_param["channel_dim"], 0) + assert_allclose(result, expected, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_as_discrete.py b/tests/test_as_discrete.py index a68e6431ec..867ef84062 100644 --- a/tests/test_as_discrete.py +++ b/tests/test_as_discrete.py @@ -66,7 +66,7 @@ class TestAsDiscrete(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_value_shape(self, input_param, img, out, expected_shape): result = AsDiscrete(**input_param)(img) - assert_allclose(result, out, rtol=1e-3) + assert_allclose(result, out, rtol=1e-3, type_test="tensor") self.assertTupleEqual(result.shape, expected_shape) diff --git a/tests/test_as_discreted.py b/tests/test_as_discreted.py index 21825c2d6c..17527c0fd4 100644 --- a/tests/test_as_discreted.py +++ b/tests/test_as_discreted.py @@ -85,10 +85,10 @@ class TestAsDiscreted(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_value_shape(self, input_param, test_input, output, expected_shape): result = AsDiscreted(**input_param)(test_input) - assert_allclose(result["pred"], output["pred"], rtol=1e-3) + assert_allclose(result["pred"], output["pred"], rtol=1e-3, type_test="tensor") self.assertTupleEqual(result["pred"].shape, expected_shape) if "label" in result: - assert_allclose(result["label"], output["label"], rtol=1e-3) + assert_allclose(result["label"], output["label"], rtol=1e-3, type_test="tensor") self.assertTupleEqual(result["label"].shape, expected_shape) diff --git a/tests/test_atss_box_matcher.py b/tests/test_atss_box_matcher.py new file mode 100644 index 0000000000..093641bb2f --- /dev/null +++ b/tests/test_atss_box_matcher.py @@ -0,0 +1,44 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.ATSS_matcher import ATSSMatcher +from monai.data.box_utils import box_iou +from tests.utils import assert_allclose + +TEST_CASES = [ + [ + {"num_candidates": 2, "similarity_fn": box_iou, "center_in_gt": False}, + torch.tensor([[0, 1, 2, 3, 2, 5]], dtype=torch.float16), + torch.tensor([[0, 1, 2, 3, 2, 5], [0, 1, 1, 3, 2, 5], [0, 1, 2, 3, 2, 4]], dtype=torch.float16), + [3], + 3, + torch.tensor([0, -1, -1]), + ] +] + + +class TestATSS(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_atss(self, input_param, boxes, anchors, num_anchors_per_level, num_anchors_per_loc, expected_matches): + matcher = ATSSMatcher(**input_param, debug=True) + match_quality_matrix, matches = matcher.compute_matches( + boxes, anchors, num_anchors_per_level, num_anchors_per_loc + ) + assert_allclose(matches, expected_matches, type_test=True, device_test=True, atol=0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bending_energy.py b/tests/test_bending_energy.py index 318b1905df..18f88ba759 100644 --- a/tests/test_bending_energy.py +++ b/tests/test_bending_energy.py @@ -60,7 +60,6 @@ def test_ill_shape(self): loss.forward(torch.ones((1, 3), device=device)) with self.assertRaisesRegex(ValueError, "Expecting 3-d, 4-d or 5-d"): loss.forward(torch.ones((1, 4, 5, 5, 5, 5), device=device)) - # spatial_dim < 5 with self.assertRaisesRegex(ValueError, "All spatial dimensions"): loss.forward(torch.ones((1, 3, 4, 5, 5), device=device)) with self.assertRaisesRegex(ValueError, "All spatial dimensions"): diff --git a/tests/test_border_pad.py b/tests/test_border_pad.py index b632ff831f..1194ae49a6 100644 --- a/tests/test_border_pad.py +++ b/tests/test_border_pad.py @@ -11,45 +11,32 @@ import unittest -import numpy as np from parameterized import parameterized from monai.transforms import BorderPad -from monai.utils import NumpyPadMode -from tests.utils import TEST_NDARRAYS - -TEST_CASE_1 = [{"spatial_border": 2, "mode": "constant"}, np.zeros((3, 8, 8, 4)), np.zeros((3, 12, 12, 8))] - -TEST_CASE_2 = [{"spatial_border": [1, 2, 3], "mode": "constant"}, np.zeros((3, 8, 8, 4)), np.zeros((3, 10, 12, 10))] - -TEST_CASE_3 = [ - {"spatial_border": [1, 2, 3, 4, 5, 6], "mode": "constant"}, - np.zeros((3, 8, 8, 4)), - np.zeros((3, 11, 15, 15)), +from monai.utils.enums import NumpyPadMode, PytorchPadMode +from tests.padders import PadTest + +TESTS = [ + [{"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)], ] -TEST_CASE_4 = [ - {"spatial_border": [1, 2, 3, 4, 5, 6], "mode": NumpyPadMode.CONSTANT}, - np.zeros((3, 8, 8, 4)), - np.zeros((3, 11, 15, 15)), -] +class TestBorderPad(PadTest): + Padder = BorderPad -class TestBorderPad(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) - def test_pad_shape(self, input_param, input_data, expected_val): - for p in TEST_NDARRAYS: - padder = BorderPad(**input_param) - r1 = padder(p(input_data)) - r2 = padder(input_data, mode=input_param["mode"]) - self.assertAlmostEqual(r1.shape, expected_val.shape) - self.assertAlmostEqual(r2.shape, expected_val.shape) + @parameterized.expand(TESTS) + def test_pad(self, input_param, input_shape, expected_shape): + modes = ["constant", NumpyPadMode.CONSTANT, PytorchPadMode.CONSTANT] + self.pad_test(input_param, input_shape, expected_shape, modes) def test_pad_kwargs(self): - padder = BorderPad(spatial_border=2, mode="constant", constant_values=((0, 0), (1, 1), (2, 2))) - result = padder(np.zeros((3, 8, 4))) - np.testing.assert_allclose(result[:, :2, 2:6], np.ones((3, 2, 4))) - np.testing.assert_allclose(result[:, :, :2], np.ones((3, 12, 2)) + 1) + kwargs = {"spatial_border": 2, "mode": "constant"} + unchanged_slices = [slice(None), slice(2, -2), slice(2, -2)] + self.pad_test_kwargs(unchanged_slices, **kwargs) if __name__ == "__main__": diff --git a/tests/test_border_padd.py b/tests/test_border_padd.py index e4b8dd20ea..ca55c8b09d 100644 --- a/tests/test_border_padd.py +++ b/tests/test_border_padd.py @@ -11,49 +11,28 @@ import unittest -import numpy as np from parameterized import parameterized from monai.transforms import BorderPadd -from monai.utils import NumpyPadMode - -TEST_CASE_1 = [ - {"keys": ["img", "seg"], "spatial_border": 2, "mode": ["constant", "edge"]}, - {"img": np.zeros((3, 8, 8, 4)), "seg": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 12, 12, 8)), -] - -TEST_CASE_2 = [ - {"keys": "img", "spatial_border": [1, 2, 3], "mode": "constant"}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 10, 12, 10)), -] - -TEST_CASE_3 = [ - {"keys": "img", "spatial_border": [1, 2, 3, 4, 5, 6], "mode": "constant"}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 11, 15, 15)), -] - -TEST_CASE_4 = [ - {"keys": ["img", "seg"], "spatial_border": 2, "mode": ["constant", NumpyPadMode.EDGE]}, - {"img": np.zeros((3, 8, 8, 4)), "seg": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 12, 12, 8)), +from monai.utils.enums import NumpyPadMode, PytorchPadMode +from tests.padders import PadTest + +TESTS = [ + [{"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)], ] -TEST_CASE_5 = [ - {"keys": ["img", "seg"], "spatial_border": 2, "mode": [NumpyPadMode.CONSTANT, NumpyPadMode.EDGE]}, - {"img": np.zeros((3, 8, 8, 4)), "seg": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 12, 12, 8)), -] +class TestBorderPadd(PadTest): + Padder = BorderPadd -class TestBorderPadd(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5]) - def test_pad_shape(self, input_param, input_data, expected_val): - padder = BorderPadd(**input_param) - result = padder(input_data) - self.assertAlmostEqual(result["img"].shape, expected_val.shape) + @parameterized.expand(TESTS) + def test_pad(self, input_param, input_shape, expected_shape): + modes = ["constant", NumpyPadMode.CONSTANT, PytorchPadMode.CONSTANT, "edge", NumpyPadMode.EDGE] + self.pad_test(input_param, input_shape, expected_shape, modes) if __name__ == "__main__": diff --git a/tests/test_box_coder.py b/tests/test_box_coder.py new file mode 100644 index 0000000000..86ca7a98c2 --- /dev/null +++ b/tests/test_box_coder.py @@ -0,0 +1,41 @@ +# 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. + +import unittest + +import torch + +from monai.apps.detection.utils.box_coder import BoxCoder +from monai.transforms import CastToType +from tests.utils import assert_allclose + + +class TestBoxTransform(unittest.TestCase): + def test_value(self): + box_coder = BoxCoder(weights=[1, 1, 1, 1, 1, 1]) + test_dtype = [torch.float32, torch.float16] + for dtype in test_dtype: + gt_boxes_0 = torch.rand((10, 3)).abs() + gt_boxes_1 = gt_boxes_0 + torch.rand((10, 3)).abs() + 10 + gt_boxes = torch.cat((gt_boxes_0, gt_boxes_1), dim=1) + gt_boxes = CastToType(dtype=dtype)(gt_boxes) + + proposals_0 = (gt_boxes_0 + torch.rand(gt_boxes_0.shape)).abs() + proposals_1 = proposals_0 + torch.rand(gt_boxes_0.shape).abs() + 10 + proposals = torch.cat((proposals_0, proposals_1), dim=1) + + rel_gt_boxes = box_coder.encode_single(gt_boxes, proposals) + gt_back = box_coder.decode_single(rel_gt_boxes, proposals) + assert_allclose(gt_back, gt_boxes, type_test=True, device_test=True, atol=0.1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_box_transform.py b/tests/test_box_transform.py new file mode 100644 index 0000000000..6b0a4a2b19 --- /dev/null +++ b/tests/test_box_transform.py @@ -0,0 +1,292 @@ +# 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. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.apps.detection.transforms.box_ops import convert_mask_to_box +from monai.apps.detection.transforms.dictionary import BoxToMaskd, ConvertBoxModed, MaskToBoxd +from monai.transforms import CastToTyped +from tests.utils import TEST_NDARRAYS, assert_allclose + +TESTS_3D = [] +boxes = [[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 3, 3], [0, 1, 1, 2, 3, 4]] +labels = [1, 1, 0] +scores = [[0.2, 0.8], [0.3, 0.7], [0.6, 0.4]] +image_size = [1, 4, 6, 4] +image = np.zeros(image_size) + +for p in TEST_NDARRAYS: + TESTS_3D.append( + [ + {"box_keys": "boxes", "dst_mode": "xyzwhd"}, + {"boxes": p(boxes), "image": p(image), "labels": p(labels), "scores": p(scores)}, + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]]), + p([[0, 0, 0, 0, 0, 0], [0, 3, 0, 1, 9, 4.5], [0, 3, 1.5, 1, 9, 6]]), + p([[1, -6, -1, 1, -6, -1], [1, -3, -1, 2, 3, 3.5], [1, -3, 0.5, 2, 3, 5]]), + p([[4, 6, 4, 4, 6, 4], [2, 3, 1, 4, 5, 4], [2, 3, 0, 4, 5, 3]]), + p([[0, 1, 0, 2, 3, 3], [0, 1, 1, 2, 3, 4]]), + p([[6, 0, 0, 6, 0, 0], [3, 0, 0, 5, 2, 3], [3, 0, 1, 5, 2, 4]]), + ] + ) + +TESTS_2D = [] +boxes = [[0, 1, 2, 2], [0, 0, 1, 1]] +labels = [1, 0] +image_size = [1, 2, 2] +image = np.zeros(image_size) +for p in TEST_NDARRAYS: + TESTS_2D.append( + [{"boxes": p(boxes), "image": p(image), "labels": p(labels)}, p([[[0, 2], [0, 2]], [[1, 0], [0, 0]]])] + ) + +TESTS_2D_mask = [] +boxes_mask = [[[-1, 0], [0, -1]]] +for p in TEST_NDARRAYS: + TESTS_2D_mask.append([p(boxes_mask), (p([[0.0, 0.0, 2.0, 2.0]]), p([0]))]) +boxes_mask = [[[-1, 0], [0, -1]], [[-1, 1], [1, -1]]] +for p in TEST_NDARRAYS: + TESTS_2D_mask.append([p(boxes_mask), (p([[0.0, 0.0, 2.0, 2.0], [0.0, 0.0, 2.0, 2.0]]), p([0, 1]))]) + + +class TestBoxTransform(unittest.TestCase): + @parameterized.expand(TESTS_2D_mask) + def test_value_2d_mask(self, mask, expected_box_label): + box_label = convert_mask_to_box(mask) + assert_allclose(box_label[0], expected_box_label[0], type_test=True, device_test=True, atol=1e-3) + assert_allclose(box_label[1], expected_box_label[1], type_test=True, device_test=True, atol=1e-3) + + @parameterized.expand(TESTS_2D) + def test_value_2d(self, data, expected_mask): + test_dtype = [torch.float32, torch.float16] + for dtype in test_dtype: + data = CastToTyped(keys=["image", "boxes"], dtype=dtype)(data) + transform_to_mask = BoxToMaskd( + box_keys="boxes", + box_mask_keys="box_mask", + box_ref_image_keys="image", + label_keys="labels", + min_fg_label=0, + ellipse_mask=False, + ) + transform_to_box = MaskToBoxd( + box_keys="boxes", box_mask_keys="box_mask", label_keys="labels", min_fg_label=0 + ) + data_mask = transform_to_mask(data) + assert_allclose(data_mask["box_mask"], expected_mask, type_test=True, device_test=True, atol=1e-3) + data_back = transform_to_box(data_mask) + assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + assert_allclose(data_back["labels"], data["labels"], type_test=False, device_test=False, atol=1e-3) + + def test_value_3d_mask(self): + test_dtype = [torch.float32, torch.float16] + image = np.zeros((1, 32, 33, 34)) + boxes = np.array([[7, 8, 9, 10, 12, 13], [1, 3, 5, 2, 5, 9], [0, 0, 0, 1, 1, 1]]) + data = {"image": image, "boxes": boxes, "labels": np.array((1, 0, 3))} + for dtype in test_dtype: + data = CastToTyped(keys=["image", "boxes"], dtype=dtype)(data) + transform_to_mask = BoxToMaskd( + box_keys="boxes", + box_mask_keys="box_mask", + box_ref_image_keys="image", + label_keys="labels", + min_fg_label=0, + ellipse_mask=False, + ) + transform_to_box = MaskToBoxd( + box_keys="boxes", box_mask_keys="box_mask", label_keys="labels", min_fg_label=0 + ) + data_mask = transform_to_mask(data) + assert_allclose(data_mask["box_mask"].shape, (3, 32, 33, 34), type_test=True, device_test=True, atol=1e-3) + data_back = transform_to_box(data_mask) + assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + assert_allclose(data_back["labels"], data["labels"], type_test=False, device_test=False, atol=1e-3) + + @parameterized.expand(TESTS_3D) + def test_value_3d( + self, + keys, + data, + expected_convert_result, + expected_zoom_result, + expected_zoom_keepsize_result, + expected_flip_result, + expected_clip_result, + expected_rotate_result, + ): + test_dtype = [torch.float32] + for dtype in test_dtype: + data = CastToTyped(keys=["image", "boxes"], dtype=dtype)(data) + # test ConvertBoxToStandardModed + transform_convert_mode = ConvertBoxModed(**keys) + convert_result = transform_convert_mode(data) + assert_allclose( + convert_result["boxes"], expected_convert_result, type_test=True, device_test=True, atol=1e-3 + ) + + # invert_transform_convert_mode = Invertd( + # keys=["boxes"], transform=transform_convert_mode, orig_keys=["boxes"] + # ) + # data_back = invert_transform_convert_mode(convert_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # + # # test ZoomBoxd + # transform_zoom = ZoomBoxd( + # image_keys="image", box_keys="boxes", box_ref_image_keys="image", zoom=[0.5, 3, 1.5], keep_size=False + # ) + # zoom_result = transform_zoom(data) + # assert_allclose(zoom_result["boxes"], expected_zoom_result, type_test=True, device_test=True, atol=1e-3) + # invert_transform_zoom = Invertd( + # keys=["image", "boxes"], transform=transform_zoom, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_zoom(zoom_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + # + # transform_zoom = ZoomBoxd( + # image_keys="image", box_keys="boxes", box_ref_image_keys="image", zoom=[0.5, 3, 1.5], keep_size=True + # ) + # zoom_result = transform_zoom(data) + # assert_allclose( + # zoom_result["boxes"], expected_zoom_keepsize_result, type_test=True, device_test=True, atol=1e-3 + # ) + # + # # test RandZoomBoxd + # transform_zoom = RandZoomBoxd( + # image_keys="image", + # box_keys="boxes", + # box_ref_image_keys="image", + # prob=1.0, + # min_zoom=(0.3,) * 3, + # max_zoom=(3.0,) * 3, + # keep_size=False, + # ) + # zoom_result = transform_zoom(data) + # invert_transform_zoom = Invertd( + # keys=["image", "boxes"], transform=transform_zoom, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_zoom(zoom_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=0.01) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + # + # # test AffineBoxToImageCoordinated, AffineBoxToWorldCoordinated + # transform_affine = AffineBoxToImageCoordinated(box_keys="boxes", box_ref_image_keys="image") + # with self.assertRaises(Exception) as context: + # transform_affine(data) + # self.assertTrue("Please check whether it is the correct the image meta key." in str(context.exception)) + # + # data["image_meta_dict"] = {"affine": torch.diag(1.0 / torch.Tensor([0.5, 3, 1.5, 1]))} + # affine_result = transform_affine(data) + # assert_allclose(affine_result["boxes"], expected_zoom_result, type_test=True, device_test=True, atol=0.01) + # invert_transform_affine = Invertd(keys=["boxes"], transform=transform_affine, orig_keys=["boxes"]) + # data_back = invert_transform_affine(affine_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=0.01) + # invert_transform_affine = AffineBoxToWorldCoordinated(box_keys="boxes", box_ref_image_keys="image") + # data_back = invert_transform_affine(affine_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=0.01) + # + # # test FlipBoxd + # transform_flip = FlipBoxd( + # image_keys="image", box_keys="boxes", box_ref_image_keys="image", spatial_axis=[0, 1, 2] + # ) + # flip_result = transform_flip(data) + # assert_allclose(flip_result["boxes"], expected_flip_result, type_test=True, device_test=True, atol=1e-3) + # invert_transform_flip = Invertd( + # keys=["image", "boxes"], transform=transform_flip, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_flip(flip_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + # + # # test RandFlipBoxd + # for spatial_axis in [(0,), (1,), (2,), (0, 1), (1, 2)]: + # transform_flip = RandFlipBoxd( + # image_keys="image", + # box_keys="boxes", + # box_ref_image_keys="image", + # prob=1.0, + # spatial_axis=spatial_axis, + # ) + # flip_result = transform_flip(data) + # invert_transform_flip = Invertd( + # keys=["image", "boxes"], transform=transform_flip, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_flip(flip_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + # + # # test ClipBoxToImaged + # transform_clip = ClipBoxToImaged( + # box_keys="boxes", box_ref_image_keys="image", label_keys=["labels", "scores"], remove_empty=True + # ) + # clip_result = transform_clip(data) + # assert_allclose(clip_result["boxes"], expected_clip_result, type_test=True, device_test=True, atol=1e-3) + # assert_allclose(clip_result["labels"], data["labels"][1:], type_test=True, device_test=True, atol=1e-3) + # assert_allclose(clip_result["scores"], data["scores"][1:], type_test=True, device_test=True, atol=1e-3) + # + # transform_clip = ClipBoxToImaged( + # box_keys="boxes", box_ref_image_keys="image", label_keys=[], remove_empty=True + # ) # corner case when label_keys is empty + # clip_result = transform_clip(data) + # assert_allclose(clip_result["boxes"], expected_clip_result, type_test=True, device_test=True, atol=1e-3) + # + # # test RandCropBoxByPosNegLabeld + # transform_crop = RandCropBoxByPosNegLabeld( + # image_keys="image", box_keys="boxes", label_keys=["labels", "scores"], spatial_size=2, num_samples=3 + # ) + # crop_result = transform_crop(data) + # assert len(crop_result) == 3 + # for ll in range(3): + # assert_allclose( + # crop_result[ll]["boxes"].shape[0], + # crop_result[ll]["labels"].shape[0], + # type_test=True, + # device_test=True, + # atol=1e-3, + # ) + # assert_allclose( + # crop_result[ll]["boxes"].shape[0], + # crop_result[ll]["scores"].shape[0], + # type_test=True, + # device_test=True, + # atol=1e-3, + # ) + # + # # test RotateBox90d + # transform_rotate = RotateBox90d( + # image_keys="image", box_keys="boxes", box_ref_image_keys="image", k=1, spatial_axes=[0, 1] + # ) + # rotate_result = transform_rotate(data) + # assert_allclose(rotate_result["boxes"], expected_rotate_result, type_test=True, device_test=True, atol=1e-3) + # invert_transform_rotate = Invertd( + # keys=["image", "boxes"], transform=transform_rotate, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_rotate(rotate_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + # + # transform_rotate = RandRotateBox90d( + # image_keys="image", box_keys="boxes", box_ref_image_keys="image", prob=1.0, max_k=3, spatial_axes=[0, 1] + # ) + # rotate_result = transform_rotate(data) + # invert_transform_rotate = Invertd( + # keys=["image", "boxes"], transform=transform_rotate, orig_keys=["image", "boxes"] + # ) + # data_back = invert_transform_rotate(rotate_result) + # assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + # assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_box_utils.py b/tests/test_box_utils.py new file mode 100644 index 0000000000..8c56783c3b --- /dev/null +++ b/tests/test_box_utils.py @@ -0,0 +1,220 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.data.box_utils import ( + CenterSizeMode, + CornerCornerModeTypeA, + CornerCornerModeTypeB, + CornerCornerModeTypeC, + CornerSizeMode, + box_area, + box_centers, + box_giou, + box_iou, + box_pair_giou, + boxes_center_distance, + centers_in_boxes, + clip_boxes_to_image, + convert_box_mode, + convert_box_to_standard_mode, + non_max_suppression, +) +from monai.utils.type_conversion import convert_data_type +from tests.utils import TEST_NDARRAYS, assert_allclose + +TESTS = [] +for p in TEST_NDARRAYS: + boxes = [[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]] + spatial_size = [4, 4, 4] + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "cccwhd", "half": False}, + CornerSizeMode, + p([[0, 0, 0, 0, 0, 0], [-1, 0, -1.5, 2, 2, 3], [-1, 0, -0.5, 2, 2, 3]]), + p([0, 12, 12]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xyzwhd", "half": False}, + CornerSizeMode, + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]]), + p([0, 12, 12]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xyzwhd", "half": True}, + "xyzxyz", + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 3, 3], [0, 1, 1, 2, 3, 4]]), + p([0, 12, 12]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xyzwhd", "half": False}, + "xxyyzz", + p([[0, 0, 0, 0, 0, 0], [0, 2, 1, 3, 0, 3], [0, 2, 1, 3, 1, 4]]), + p([0, 12, 12]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xyzwhd", "half": False}, + CornerCornerModeTypeC, + p([[0, 0, 0, 0, 0, 0], [0, 1, 2, 3, 0, 3], [0, 1, 2, 3, 1, 4]]), + p([0, 12, 12]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": CornerCornerModeTypeA(), "half": False}, + "xyzwhd", + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 1, 3], [0, 1, 1, 2, 1, 2]]), + p([0, 6, 4]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": CornerCornerModeTypeA, "half": True}, + CornerCornerModeTypeA, + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]]), + p([0, 6, 4]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xyzxyz", "half": False}, + CornerCornerModeTypeB(), + p([[0, 0, 0, 0, 0, 0], [0, 2, 1, 2, 0, 3], [0, 2, 1, 2, 1, 3]]), + p([0, 6, 4]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xxyyzz", "half": False}, + "xxyyzz", + p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]]), + p([0, 2, 1]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xxyyzz", "half": True}, + "xyzxyz", + p([[0, 0, 0, 0, 0, 0], [0, 0, 2, 1, 2, 3], [0, 1, 2, 1, 2, 3]]), + p([0, 2, 1]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xxyyzz", "half": False}, + "xyzwhd", + p([[0, 0, 0, 0, 0, 0], [0, 0, 2, 1, 2, 1], [0, 1, 2, 1, 1, 1]]), + p([0, 2, 1]), + ] + ) + TESTS.append( + [ + {"boxes": p(boxes), "spatial_size": spatial_size, "mode": "xxyyzz", "half": False}, + CenterSizeMode(), + p([[0, 0, 0, 0, 0, 0], [0.5, 1, 2.5, 1, 2, 1], [0.5, 1.5, 2.5, 1, 1, 1]]), + p([0, 2, 1]), + ] + ) + + +class TestCreateBoxList(unittest.TestCase): + @parameterized.expand(TESTS) + def test_value(self, input_data, mode2, expected_box, expected_area): + expected_box = convert_data_type(expected_box, dtype=np.float32)[0] + boxes1 = convert_data_type(input_data["boxes"], dtype=np.float32)[0] + mode1 = input_data["mode"] + half_bool = input_data["half"] + spatial_size = input_data["spatial_size"] + + # test float16 + if half_bool: + boxes1 = convert_data_type(boxes1, dtype=np.float16)[0] + expected_box = convert_data_type(expected_box, dtype=np.float16)[0] + + # test convert_box_mode, convert_box_to_standard_mode + result2 = convert_box_mode(boxes=boxes1, src_mode=mode1, dst_mode=mode2) + assert_allclose(result2, expected_box, type_test=True, device_test=True, atol=0.0) + + result1 = convert_box_mode(boxes=result2, src_mode=mode2, dst_mode=mode1) + assert_allclose(result1, boxes1, type_test=True, device_test=True, atol=0.0) + + result_standard = convert_box_to_standard_mode(boxes=boxes1, mode=mode1) + expected_box_standard = convert_box_to_standard_mode(boxes=expected_box, mode=mode2) + assert_allclose(result_standard, expected_box_standard, type_test=True, device_test=True, atol=0.0) + + # test box_area, box_iou, box_giou, box_pair_giou + assert_allclose(box_area(result_standard), expected_area, type_test=True, device_test=True, atol=0.0) + iou_metrics = (box_iou, box_giou) + for p in iou_metrics: + self_iou = p(boxes1=result_standard[1:2, :], boxes2=result_standard[1:1, :]) + assert_allclose(self_iou, np.array([[]]), type_test=False) + + self_iou = p(boxes1=result_standard[1:2, :], boxes2=result_standard[1:2, :]) + assert_allclose(self_iou, np.array([[1.0]]), type_test=False) + + self_iou = box_pair_giou(boxes1=result_standard[1:1, :], boxes2=result_standard[1:1, :]) + assert_allclose(self_iou, np.array([]), type_test=False) + + self_iou = box_pair_giou(boxes1=result_standard[1:2, :], boxes2=result_standard[1:2, :]) + assert_allclose(self_iou, np.array([1.0]), type_test=False) + + # test box_centers, centers_in_boxes, boxes_center_distance + result_standard_center = box_centers(result_standard) + expected_center = convert_box_mode(boxes=boxes1, src_mode=mode1, dst_mode="cccwhd")[:, :3] + assert_allclose(result_standard_center, expected_center, type_test=True, device_test=True, atol=0.0) + + center = expected_center + center[2, :] += 10 + result_centers_in_boxes = centers_in_boxes(centers=center, boxes=result_standard) + assert_allclose(result_centers_in_boxes, np.array([False, True, False]), type_test=False) + + center_dist, _, _ = boxes_center_distance(boxes1=result_standard[1:2, :], boxes2=result_standard[1:1, :]) + assert_allclose(center_dist, np.array([[]]), type_test=False) + center_dist, _, _ = boxes_center_distance(boxes1=result_standard[1:2, :], boxes2=result_standard[1:2, :]) + assert_allclose(center_dist, np.array([[0.0]]), type_test=False) + center_dist, _, _ = boxes_center_distance(boxes1=result_standard[0:1, :], boxes2=result_standard[0:1, :]) + assert_allclose(center_dist, np.array([[0.0]]), type_test=False) + + # test clip_boxes_to_image + clipped_boxes, keep = clip_boxes_to_image(expected_box_standard, spatial_size, remove_empty=True) + assert_allclose( + expected_box_standard[keep, :], expected_box_standard[1:, :], type_test=True, device_test=True, atol=0.0 + ) + assert_allclose( + id(clipped_boxes) != id(expected_box_standard), True, type_test=False, device_test=False, atol=0.0 + ) + + # test non_max_suppression + nms_box = non_max_suppression( + boxes=result_standard, scores=boxes1[:, 1] / 2.0, nms_thresh=1.0, box_overlap_metric=box_giou + ) + assert_allclose(nms_box, [1, 2, 0], type_test=False) + + nms_box = non_max_suppression( + boxes=result_standard, scores=boxes1[:, 1] / 2.0, nms_thresh=-1.0, box_overlap_metric=box_iou + ) + assert_allclose(nms_box, [1], type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bundle_ckpt_export.py b/tests/test_bundle_ckpt_export.py index 0f7d0f7d35..a7cbff22f0 100644 --- a/tests/test_bundle_ckpt_export.py +++ b/tests/test_bundle_ckpt_export.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import subprocess import tempfile @@ -17,6 +18,7 @@ from parameterized import parameterized from monai.bundle import ConfigParser +from monai.data import load_net_with_metadata from monai.networks import save_state from tests.utils import skip_if_windows @@ -33,7 +35,8 @@ def test_export(self, key_in_ckpt): config_file = os.path.join(os.path.dirname(__file__), "testing_data", "inference.json") with tempfile.TemporaryDirectory() as tempdir: def_args = {"meta_file": "will be replaced by `meta_file` arg"} - def_args_file = os.path.join(tempdir, "def_args.json") + def_args_file = os.path.join(tempdir, "def_args.yaml") + ckpt_file = os.path.join(tempdir, "model.pt") ts_file = os.path.join(tempdir, "model.ts") @@ -44,11 +47,18 @@ def test_export(self, key_in_ckpt): save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file) cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "network_def", "--filepath", ts_file] - cmd += ["--meta_file", meta_file, "--config_file", config_file, "--ckpt_file", ckpt_file] - cmd += ["--key_in_ckpt", key_in_ckpt, "--args_file", def_args_file] + cmd += ["--meta_file", meta_file, "--config_file", f"['{config_file}','{def_args_file}']", "--ckpt_file"] + cmd += [ckpt_file, "--key_in_ckpt", key_in_ckpt, "--args_file", def_args_file] subprocess.check_call(cmd) self.assertTrue(os.path.exists(ts_file)) + _, metadata, extra_files = load_net_with_metadata( + ts_file, more_extra_files=["inference.json", "def_args.json"] + ) + self.assertTrue("schema" in metadata) + self.assertTrue("meta_file" in json.loads(extra_files["def_args.json"])) + self.assertTrue("network_def" in json.loads(extra_files["inference.json"])) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_bundle_download.py b/tests/test_bundle_download.py new file mode 100644 index 0000000000..7e609a7b31 --- /dev/null +++ b/tests/test_bundle_download.py @@ -0,0 +1,173 @@ +# 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. + +import json +import os +import subprocess +import tempfile +import unittest + +import torch +from parameterized import parameterized + +import monai.networks.nets as nets +from monai.apps import check_hash +from monai.bundle import ConfigParser, load +from tests.utils import SkipIfBeforePyTorchVersion, skip_if_downloading_fails, skip_if_quick, skip_if_windows + +TEST_CASE_1 = [ + ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"], + "test_bundle", + "Project-MONAI/MONAI-extra-test-data/0.8.1", + "a131d39a0af717af32d19e565b434928", +] + +TEST_CASE_2 = [ + ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"], + "test_bundle", + "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/test_bundle.zip", + "a131d39a0af717af32d19e565b434928", +] + +TEST_CASE_3 = [ + ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"], + "test_bundle", + "Project-MONAI/MONAI-extra-test-data/0.8.1", + "cuda" if torch.cuda.is_available() else "cpu", + "model.pt", +] + +TEST_CASE_4 = [ + ["test_output.pt", "test_input.pt"], + "test_bundle", + "Project-MONAI/MONAI-extra-test-data/0.8.1", + "cuda" if torch.cuda.is_available() else "cpu", + "model.ts", +] + + +@skip_if_windows +class TestDownload(unittest.TestCase): + @parameterized.expand([TEST_CASE_1]) + @skip_if_quick + def test_download_bundle(self, bundle_files, bundle_name, repo, hash_val): + with skip_if_downloading_fails(): + # download a whole bundle from github releases + with tempfile.TemporaryDirectory() as tempdir: + cmd = ["coverage", "run", "-m", "monai.bundle", "download", "--name", bundle_name, "--source", "github"] + cmd += ["--bundle_dir", tempdir, "--repo", repo, "--progress", "False"] + subprocess.check_call(cmd) + 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)) + + @parameterized.expand([TEST_CASE_2]) + @skip_if_quick + def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val): + with skip_if_downloading_fails(): + # download a single file from url, also use `args_file` + with tempfile.TemporaryDirectory() as tempdir: + def_args = {"name": bundle_name, "bundle_dir": tempdir, "url": ""} + def_args_file = os.path.join(tempdir, "def_args.json") + parser = ConfigParser() + parser.export_config_file(config=def_args, filepath=def_args_file) + cmd = ["coverage", "run", "-m", "monai.bundle", "download", "--args_file", def_args_file] + cmd += ["--url", url] + subprocess.check_call(cmd) + 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)) + + +class TestLoad(unittest.TestCase): + @parameterized.expand([TEST_CASE_3]) + @skip_if_quick + def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file): + with skip_if_downloading_fails(): + # download bundle, and load weights from the downloaded path + with tempfile.TemporaryDirectory() as tempdir: + # load weights + weights = load( + name=bundle_name, + model_file=model_file, + bundle_dir=tempdir, + repo=repo, + progress=False, + device=device, + ) + + # prepare network + with open(os.path.join(tempdir, bundle_name, bundle_files[2])) as f: + net_args = json.load(f)["network_def"] + model_name = net_args["_target_"] + del net_args["_target_"] + model = nets.__dict__[model_name](**net_args) + model.to(device) + model.load_state_dict(weights) + model.eval() + + # prepare data and test + input_tensor = torch.load(os.path.join(tempdir, bundle_name, bundle_files[4]), map_location=device) + output = model.forward(input_tensor) + expected_output = torch.load(os.path.join(tempdir, bundle_name, bundle_files[3]), map_location=device) + torch.testing.assert_allclose(output, expected_output) + + # load instantiated model directly and test, since the bundle has been downloaded, + # there is no need to input `repo` + model_2 = load( + name=bundle_name, + model_file=model_file, + bundle_dir=tempdir, + progress=False, + device=device, + net_name=model_name, + **net_args, + ) + model_2.eval() + output_2 = model_2.forward(input_tensor) + torch.testing.assert_allclose(output_2, expected_output) + + @parameterized.expand([TEST_CASE_4]) + @skip_if_quick + @SkipIfBeforePyTorchVersion((1, 7, 1)) + def test_load_ts_module(self, bundle_files, bundle_name, repo, device, model_file): + with skip_if_downloading_fails(): + # load ts module + with tempfile.TemporaryDirectory() as tempdir: + # load ts module + model_ts, metadata, extra_file_dict = load( + name=bundle_name, + model_file=model_file, + load_ts_module=True, + bundle_dir=tempdir, + repo=repo, + progress=False, + device=device, + config_files=("network.json",), + ) + + # prepare and test ts + input_tensor = torch.load(os.path.join(tempdir, bundle_name, bundle_files[1]), map_location=device) + output = model_ts.forward(input_tensor) + expected_output = torch.load(os.path.join(tempdir, bundle_name, bundle_files[0]), map_location=device) + torch.testing.assert_allclose(output, expected_output) + # test metadata + self.assertTrue(metadata["pytorch_version"] == "1.7.1") + # test extra_file_dict + self.assertTrue("network.json" in extra_file_dict.keys()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bundle_init_bundle.py b/tests/test_bundle_init_bundle.py new file mode 100644 index 0000000000..24fc425c31 --- /dev/null +++ b/tests/test_bundle_init_bundle.py @@ -0,0 +1,41 @@ +# 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. + +import os +import subprocess +import tempfile +import unittest + +import torch + +from monai.networks.nets import UNet +from tests.utils import skip_if_windows + + +@skip_if_windows +class TestBundleInit(unittest.TestCase): + def test_bundle(self): + with tempfile.TemporaryDirectory() as tempdir: + net = UNet(2, 1, 1, [4, 8], [2]) + torch.save(net.state_dict(), tempdir + "/test.pt") + + bundle_root = tempdir + "/test_bundle" + + cmd = ["coverage", "run", "-m", "monai.bundle", "init_bundle", bundle_root, tempdir + "/test.pt"] + subprocess.check_call(cmd) + + self.assertTrue(os.path.exists(bundle_root + "/configs/metadata.json")) + self.assertTrue(os.path.exists(bundle_root + "/configs/inference.json")) + self.assertTrue(os.path.exists(bundle_root + "/models/model.pt")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bundle_utils.py b/tests/test_bundle_utils.py new file mode 100644 index 0000000000..46b29651cd --- /dev/null +++ b/tests/test_bundle_utils.py @@ -0,0 +1,117 @@ +# 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. + +import os +import shutil +import subprocess +import tempfile +import unittest + +import torch + +from monai.bundle.utils import load_bundle_config +from monai.networks.nets import UNet +from tests.utils import skip_if_windows + +metadata = """ +{ + "test_value": 1, + "test_list": [2,3] +} +""" + +test_json = """ +{ + "test_dict": { + "a": 3, + "b": "c" + }, + "network_def": { + "_target_": "UNet", + "spatial_dims": 2, + "in_channels": 1, + "out_channels": 1, + "channels": [4,8], + "strides": [2] + } +} +""" + + +@skip_if_windows +class TestLoadBundleConfig(unittest.TestCase): + def setUp(self): + self.bundle_dir = tempfile.TemporaryDirectory() + self.dir_name = os.path.join(self.bundle_dir.name, "TestBundle") + self.configs_name = os.path.join(self.dir_name, "configs") + self.models_name = os.path.join(self.dir_name, "models") + self.metadata_name = os.path.join(self.configs_name, "metadata.json") + self.test_name = os.path.join(self.configs_name, "test.json") + self.modelpt_name = os.path.join(self.models_name, "model.pt") + + self.zip_file = os.path.join(self.bundle_dir.name, "TestBundle.zip") + self.ts_file = os.path.join(self.bundle_dir.name, "TestBundle.ts") + + # create the directories for the bundle + os.mkdir(self.dir_name) + os.mkdir(self.configs_name) + os.mkdir(self.models_name) + + # fill bundle configs + + with open(self.metadata_name, "w") as o: + o.write(metadata) + + with open(self.test_name, "w") as o: + o.write(test_json) + + # save network + net = UNet(2, 1, 1, [4, 8], [2]) + torch.save(net.state_dict(), self.modelpt_name) + + def tearDown(self): + self.bundle_dir.cleanup() + + def test_load_config_dir(self): + p = load_bundle_config(self.dir_name, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + def test_load_config_zip(self): + # create a zip of the bundle + shutil.make_archive(self.zip_file[:-4], "zip", self.bundle_dir.name) + + p = load_bundle_config(self.zip_file, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + def test_load_config_ts(self): + # create a Torchscript zip of the bundle + cmd = ["python", "-m", "monai.bundle", "ckpt_export", "network_def", "--filepath", self.ts_file] + cmd += ["--meta_file", self.metadata_name] + cmd += ["--config_file", self.test_name] + cmd += ["--ckpt_file", self.modelpt_name] + + subprocess.check_output(cmd, stderr=subprocess.STDOUT) + + p = load_bundle_config(self.ts_file, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bundle_verify_net.py b/tests/test_bundle_verify_net.py index 33d480d83f..05a0731b43 100644 --- a/tests/test_bundle_verify_net.py +++ b/tests/test_bundle_verify_net.py @@ -35,8 +35,8 @@ def test_verify(self, meta_file, config_file): ConfigParser.export_config_file(config=def_args, filepath=def_args_file) cmd = ["coverage", "run", "-m", "monai.bundle", "verify_net_in_out", "network_def", "--meta_file"] - cmd += [meta_file, "--config_file", config_file, "-n", "2", "--any", "32", "--args_file", def_args_file] - cmd += ["--_meta_#network_data_format#inputs#image#spatial_shape", "[32,'*','4**p*n']"] + cmd += [meta_file, "--config_file", config_file, "-n", "4", "--any", "16", "--args_file", def_args_file] + cmd += ["--_meta_#network_data_format#inputs#image#spatial_shape", "[16,'*','2**p*n']"] test_env = os.environ.copy() print(f"CUDA_VISIBLE_DEVICES in {__file__}", test_env.get("CUDA_VISIBLE_DEVICES")) diff --git a/tests/test_cachedataset.py b/tests/test_cachedataset.py index 7227f53e04..4fa1b5ea69 100644 --- a/tests/test_cachedataset.py +++ b/tests/test_cachedataset.py @@ -55,6 +55,12 @@ def test_shape(self, transform, expected_shape): data4 = dataset[-1] self.assertEqual(len(data3), 1) + if transform is None: + # Check without providing transfrom + dataset2 = CacheDataset(data=test_data, cache_rate=0.5, as_contiguous=True) + for k in ["image", "label", "extra"]: + self.assertEqual(dataset[0][k], dataset2[0][k]) + if transform is None: self.assertEqual(data1["image"], os.path.join(tempdir, "image1.nii.gz")) self.assertEqual(data2["label"], os.path.join(tempdir, "label2.nii.gz")) @@ -80,7 +86,7 @@ def test_set_data(self): cache_rate=1.0, num_workers=4, progress=True, - copy_cache=False if sys.platform == "linux" else True, + copy_cache=not sys.platform == "linux", ) num_workers = 2 if sys.platform == "linux" else 0 diff --git a/tests/test_cachedataset_persistent_workers.py b/tests/test_cachedataset_persistent_workers.py index 4bea0486bc..7f241899eb 100644 --- a/tests/test_cachedataset_persistent_workers.py +++ b/tests/test_cachedataset_persistent_workers.py @@ -13,10 +13,8 @@ from monai.data import CacheDataset, DataLoader, create_test_image_2d from monai.transforms import Compose, RandAffined, Spacingd -from tests.utils import SkipIfBeforePyTorchVersion -@SkipIfBeforePyTorchVersion((1, 7)) class TestTransformsWCacheDatasetAndPersistentWorkers(unittest.TestCase): def test_duplicate_transforms(self): data = [{"img": create_test_image_2d(128, 128, num_seg_classes=1, channel_dim=0)[0]} for _ in range(2)] @@ -33,7 +31,7 @@ def test_duplicate_transforms(self): b1 = next(iter(train_loader)) b2 = next(iter(train_loader)) - self.assertEqual(len(b1["img_transforms"]), len(b2["img_transforms"])) + self.assertEqual(len(b1["img"].applied_operations), len(b2["img"].applied_operations)) if __name__ == "__main__": diff --git a/tests/test_center_scale_crop.py b/tests/test_center_scale_crop.py index f22651e3e0..ab07a44eb5 100644 --- a/tests/test_center_scale_crop.py +++ b/tests/test_center_scale_crop.py @@ -9,43 +9,40 @@ # See the License for the specific language governing permissions and # limitations under the License. + import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import CenterScaleCrop +from tests.croppers import CropTest -TEST_CASE_0 = [{"roi_scale": [0.6, 0.3, -1]}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 1, 3)] - -TEST_CASE_1 = [{"roi_scale": 0.6}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 2, 2)] - -TEST_CASE_2 = [ - {"roi_scale": [0.4, 0.4]}, - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), - np.array([[[1, 2], [2, 3]]]), +TEST_SHAPES = [ + [{"roi_scale": [0.6, 0.3, -1]}, (3, 3, 3, 3), (3, 2, 1, 3)], + [{"roi_scale": 0.6}, (3, 3, 3, 3), (3, 2, 2, 2)], + [{"roi_scale": 0.5}, (3, 3, 3, 3), (3, 2, 2, 2)], ] -TEST_CASE_3 = [ - {"roi_scale": 0.5}, - torch.randint(0, 2, size=[3, 3, 3, 3], device="cuda" if torch.cuda.is_available() else "cpu"), - (3, 2, 2, 2), +TEST_VALUES = [ + [ + {"roi_scale": [0.4, 0.4]}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + np.array([[[1, 2], [2, 3]]]), + ] ] -class TestCenterScaleCrop(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_3]) - def test_shape(self, input_param, input_data, expected_shape): - result = CenterScaleCrop(**input_param)(input_data) - self.assertEqual(isinstance(result, torch.Tensor), isinstance(input_data, torch.Tensor)) - np.testing.assert_allclose(result.shape, expected_shape) +class TestCenterSpatialCrop(CropTest): + Cropper = CenterScaleCrop + + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_2]) - def test_value(self, input_param, input_data, expected_value): - result = CenterScaleCrop(**input_param)(input_data) - self.assertEqual(isinstance(result, torch.Tensor), isinstance(input_data, torch.Tensor)) - np.testing.assert_allclose(result, expected_value) + @parameterized.expand(TEST_VALUES) + def test_value(self, input_param, input_arr, expected_arr): + self.crop_test_value(input_param, input_arr, expected_arr) if __name__ == "__main__": diff --git a/tests/test_center_scale_cropd.py b/tests/test_center_scale_cropd.py index 8aef2dbe5b..894692530d 100644 --- a/tests/test_center_scale_cropd.py +++ b/tests/test_center_scale_cropd.py @@ -12,44 +12,37 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import CenterScaleCropd +from tests.croppers import CropTest -TEST_CASE_0 = [{"keys": "img", "roi_scale": [0.6, 0.3, -1]}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 1, 3)] - -TEST_CASE_1 = [{"keys": "img", "roi_scale": 0.6}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 2, 2)] - -TEST_CASE_2 = [ - {"keys": "img", "roi_scale": [0.4, 0.4]}, - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), - np.array([[[1, 2], [2, 3]]]), +TESTS = [ + [{"keys": "img", "roi_scale": [0.6, 0.3, -1]}, (3, 3, 3, 3), (3, 2, 1, 3)], + [{"keys": "img", "roi_scale": 0.6}, (3, 3, 3, 3), (3, 2, 2, 2)], + [{"keys": "img", "roi_scale": 0.5}, (3, 3, 3, 3), (3, 2, 2, 2)], ] -TEST_CASE_3 = [ - {"keys": "img", "roi_scale": 0.5}, - torch.randint(0, 2, size=[3, 3, 3, 3], device="cuda" if torch.cuda.is_available() else "cpu"), - (3, 2, 2, 2), -] -TEST_CASE_4 = [ - {"keys": "test", "roi_scale": 0.6, "allow_missing_keys": True}, - np.random.randint(0, 2, size=[3, 3, 3, 3]), - (3, 3, 3, 3), +TEST_VALUES = [ + [ + {"keys": "img", "roi_scale": [0.4, 0.4]}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + np.array([[[1, 2], [2, 3]]]), + ] ] -class TestCenterScaleCropd(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_3, TEST_CASE_4]) - def test_shape(self, input_param, input_data, expected_shape): - result = CenterScaleCropd(**input_param)({"img": input_data}) - np.testing.assert_allclose(result["img"].shape, expected_shape) +class TestCenterScaleCropd(CropTest): + Cropper = CenterScaleCropd + + @parameterized.expand(TESTS) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_2]) - def test_value(self, input_param, input_data, expected_value): - result = CenterScaleCropd(**input_param)({"img": input_data}) - np.testing.assert_allclose(result["img"], expected_value) + @parameterized.expand(TEST_VALUES) + def test_value(self, input_param, input_arr, expected_arr): + self.crop_test_value(input_param, input_arr, expected_arr) if __name__ == "__main__": diff --git a/tests/test_center_spatial_crop.py b/tests/test_center_spatial_crop.py index 09f61be2f1..7b5b19107d 100644 --- a/tests/test_center_spatial_crop.py +++ b/tests/test_center_spatial_crop.py @@ -12,40 +12,36 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import CenterSpatialCrop +from tests.croppers import CropTest -TEST_CASE_0 = [{"roi_size": [2, 2, -1]}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 2, 3)] - -TEST_CASE_1 = [{"roi_size": [2, 2, 2]}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 2, 2, 2)] - -TEST_CASE_2 = [ - {"roi_size": [2, 2]}, - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), - np.array([[[1, 2], [2, 3]]]), +TEST_SHAPES = [ + [{"roi_size": [2, 2, -1]}, (3, 3, 3, 3), (3, 2, 2, 3)], + [{"roi_size": [2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2)], + [{"roi_size": [2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2)], ] -TEST_CASE_3 = [ - {"roi_size": [2, 2, 2]}, - torch.randint(0, 2, size=[3, 3, 3, 3], device="cuda" if torch.cuda.is_available() else "cpu"), - (3, 2, 2, 2), +TEST_VALUES = [ + [ + {"roi_size": [2, 2]}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + np.array([[[1, 2], [2, 3]]]), + ] ] -class TestCenterSpatialCrop(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_3]) - def test_shape(self, input_param, input_data, expected_shape): - result = CenterSpatialCrop(**input_param)(input_data) - self.assertEqual(isinstance(result, torch.Tensor), isinstance(input_data, torch.Tensor)) - np.testing.assert_allclose(result.shape, expected_shape) +class TestCenterSpatialCrop(CropTest): + Cropper = CenterSpatialCrop + + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_2]) - def test_value(self, input_param, input_data, expected_value): - result = CenterSpatialCrop(**input_param)(input_data) - self.assertEqual(isinstance(result, torch.Tensor), isinstance(input_data, torch.Tensor)) - np.testing.assert_allclose(result, expected_value) + @parameterized.expand(TEST_VALUES) + def test_value(self, input_param, input_arr, expected_arr): + self.crop_test_value(input_param, input_arr, expected_arr) if __name__ == "__main__": diff --git a/tests/test_center_spatial_cropd.py b/tests/test_center_spatial_cropd.py index bdbc1a5031..fa7bc8c8fa 100644 --- a/tests/test_center_spatial_cropd.py +++ b/tests/test_center_spatial_cropd.py @@ -15,43 +15,42 @@ from parameterized import parameterized from monai.transforms import CenterSpatialCropd -from tests.utils import TEST_NDARRAYS, assert_allclose - -TEST_SHAPES = [] -for p in TEST_NDARRAYS: - TEST_SHAPES.append( - [{"keys": "img", "roi_size": [2, -1, -1]}, {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, (3, 2, 3, 3)] - ) - - TEST_SHAPES.append( - [{"keys": "img", "roi_size": [2, 2, 2]}, {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, (3, 2, 2, 2)] - ) - -TEST_CASES = [] -for p in TEST_NDARRAYS: - TEST_CASES.append( - [ - {"keys": "img", "roi_size": [2, 2]}, - { - "img": p( - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]) - ) - }, - p(np.array([[[1, 2], [2, 3]]])), - ] - ) - - -class TestCenterSpatialCropd(unittest.TestCase): +from tests.croppers import CropTest + +TEST_SHAPES = [ + [ + {"keys": "img", "roi_size": [2, -1, -1]}, + (3, 3, 3, 3), + (3, 2, 3, 3), + (slice(None), slice(None, -1), slice(None), slice(None)), + ], + [ + {"keys": "img", "roi_size": [2, 2, 2]}, + (3, 3, 3, 3), + (3, 2, 2, 2), + (slice(None), slice(None, -1), slice(None, -1), slice(None, -1)), + ], +] + +TEST_CASES = [ + [ + {"keys": "img", "roi_size": [2, 2]}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + np.array([[[1, 2], [2, 3]]]), + ] +] + + +class TestCenterSpatialCropd(CropTest): + Cropper = CenterSpatialCropd + @parameterized.expand(TEST_SHAPES) - def test_shape(self, input_param, input_data, expected_shape): - result = CenterSpatialCropd(**input_param)(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + def test_shape(self, input_param, input_shape, expected_shape, same_area): + self.crop_test(input_param, input_shape, expected_shape, same_area) @parameterized.expand(TEST_CASES) def test_value(self, input_param, input_data, expected_value): - result = CenterSpatialCropd(**input_param)(input_data) - assert_allclose(result["img"], expected_value, type_test=False) + self.crop_test_value(input_param, input_data, expected_value) if __name__ == "__main__": diff --git a/tests/test_compute_confusion_matrix.py b/tests/test_compute_confusion_matrix.py index 1212715548..0e38357d12 100644 --- a/tests/test_compute_confusion_matrix.py +++ b/tests/test_compute_confusion_matrix.py @@ -262,7 +262,7 @@ def test_clf_with_nan(self, input_data, expected_value): metric = ConfusionMatrixMetric(**params) result = metric(**vals) np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4) - result, _ = metric.aggregate()[0] + result, _ = metric.aggregate(reduction="mean_channel")[0] expected_value, _ = do_metric_reduction(expected_value, "mean_channel") expected_value = compute_confusion_matrix_metric("tpr", expected_value) np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4) diff --git a/tests/test_compute_generalized_dice.py b/tests/test_compute_generalized_dice.py new file mode 100644 index 0000000000..38f6e57d32 --- /dev/null +++ b/tests/test_compute_generalized_dice.py @@ -0,0 +1,154 @@ +# 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. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.metrics import GeneralizedDiceScore, compute_generalized_dice + +# keep background +TEST_CASE_1 = [ # y (1, 1, 2, 2), y_pred (1, 1, 2, 2), expected out (1) + { + "y_pred": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]), + "y": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]]), + "include_background": True, + }, + [0.8], +] + +# remove background +TEST_CASE_2 = [ # y (2, 1, 2, 2), y_pred (2, 3, 2, 2), expected out (2) (no background) + { + "y_pred": torch.tensor( + [ + [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]], + [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]], + ] + ), + "y": torch.tensor( + [ + [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]]], + [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]]], + ] + ), + "include_background": False, + }, + [0.1667, 0.6667], +] + +# should return 0 for both cases +TEST_CASE_3 = [ + { + "y_pred": torch.tensor( + [ + [[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.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]], [[1.0, 1.0], [1.0, 1.0]]], + ] + ), + "y": torch.tensor( + [ + [[[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, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]], + ] + ), + "include_background": True, + }, + [0.0, 0.0], +] + +TEST_CASE_4 = [ + {"include_background": True, "reduction": "mean_batch"}, + { + "y_pred": torch.tensor( + [ + [[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]], + [[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]], + ] + ), + "y": torch.tensor( + [ + [[[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, 0.0], [0.0, 1.0]], [[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]]], + ] + ), + }, + [0.5455], +] + +TEST_CASE_5 = [ + {"include_background": True, "reduction": "sum_batch"}, + { + "y_pred": torch.tensor( + [ + [[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]], + [[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]], + ] + ), + "y": torch.tensor( + [ + [[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.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]]], + ] + ), + }, + 1.0455, +] + +TEST_CASE_6 = [{"y": torch.ones((2, 2, 3, 3)), "y_pred": torch.ones((2, 2, 3, 3))}, [1.0000, 1.0000]] + +TEST_CASE_7 = [{"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.ones((2, 2, 3, 3))}, [0.0000, 0.0000]] + +TEST_CASE_8 = [{"y": torch.ones((2, 2, 3, 3)), "y_pred": torch.zeros((2, 2, 3, 3))}, [0.0000, 0.0000]] + +TEST_CASE_9 = [{"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.zeros((2, 2, 3, 3))}, [1.0000, 1.0000]] + + +class TestComputeMeanDice(unittest.TestCase): + # Functional part tests + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8, TEST_CASE_9]) + def test_value(self, input_data, expected_value): + result = compute_generalized_dice(**input_data) + np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) + + # Functional part tests + @parameterized.expand([TEST_CASE_3]) + def test_nans(self, input_data, expected_value): + result = compute_generalized_dice(**input_data) + self.assertTrue(np.allclose(np.isnan(result.cpu().numpy()), expected_value)) + + # Samplewise tests + @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) + def test_value_class(self, input_data, expected_value): + + # same test as for compute_meandice + vals = {} + vals["y_pred"] = input_data.pop("y_pred") + vals["y"] = input_data.pop("y") + generalized_dice_score = GeneralizedDiceScore(**input_data) + generalized_dice_score(**vals) + result = generalized_dice_score.aggregate(reduction="none") + np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) + + # Aggregation tests + @parameterized.expand([TEST_CASE_4, TEST_CASE_5]) + def test_nans_class(self, params, input_data, expected_value): + + generalized_dice_score = GeneralizedDiceScore(**params) + generalized_dice_score(**input_data) + result = generalized_dice_score.aggregate() + np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_compute_meandice.py b/tests/test_compute_meandice.py index ad66ed672a..c925c6f148 100644 --- a/tests/test_compute_meandice.py +++ b/tests/test_compute_meandice.py @@ -172,9 +172,19 @@ [[1.0000, 1.0000], [1.0000, 1.0000]], ] +TEST_CASE_11 = [ + {"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.zeros((2, 2, 3, 3)), "ignore_empty": False}, + [[1.0000, 1.0000], [1.0000, 1.0000]], +] + +TEST_CASE_12 = [ + {"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.ones((2, 2, 3, 3)), "ignore_empty": False}, + [[0.0000, 0.0000], [0.0000, 0.0000]], +] + class TestComputeMeanDice(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_9]) + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_9, TEST_CASE_11, TEST_CASE_12]) def test_value(self, input_data, expected_value): result = compute_meandice(**input_data) np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) @@ -192,9 +202,9 @@ def test_value_class(self, input_data, expected_value): vals = {} vals["y_pred"] = input_data.pop("y_pred") vals["y"] = input_data.pop("y") - dice_metric = DiceMetric(**input_data, reduction="none") + dice_metric = DiceMetric(**input_data) dice_metric(**vals) - result = dice_metric.aggregate() + result = dice_metric.aggregate(reduction="none") np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) @parameterized.expand([TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8]) diff --git a/tests/test_compute_regression_metrics.py b/tests/test_compute_regression_metrics.py index 65ca73a4ec..cab1184812 100644 --- a/tests/test_compute_regression_metrics.py +++ b/tests/test_compute_regression_metrics.py @@ -75,9 +75,9 @@ def test_shape_reduction(self): out_tensor = mt.aggregate() self.assertTrue(len(out_tensor.shape) == 0) - mt = mt_fn(reduction="mean_channel") + mt = mt_fn(reduction="sum") # test reduction arg overriding mt(in_tensor, in_tensor) - out_tensor = mt.aggregate() + out_tensor = mt.aggregate(reduction="mean_channel") self.assertTrue(len(out_tensor.shape) == 1 and out_tensor.shape[0] == batch) mt = mt_fn(reduction="sum_channel") @@ -109,9 +109,9 @@ def test_compare_numpy(self): # check metrics for mt_fn, mt_fn_np in zip(metrics, metrics_np): - mt = mt_fn(reduction="mean") + mt = mt_fn() mt(y_pred=in_tensor_a, y=in_tensor_b) - out_tensor = mt.aggregate() + out_tensor = mt.aggregate(reduction="mean") out_np = mt_fn_np(y_pred=in_tensor_a.cpu().numpy(), y=in_tensor_b.cpu().numpy()) np.testing.assert_allclose(out_tensor.cpu().numpy(), out_np, atol=1e-4) diff --git a/tests/test_compute_roc_auc.py b/tests/test_compute_roc_auc.py index 887db08c7c..2c9135024f 100644 --- a/tests/test_compute_roc_auc.py +++ b/tests/test_compute_roc_auc.py @@ -141,6 +141,8 @@ def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_valu metric = ROCAUCMetric(average=average) metric(y_pred=y_pred, y=y) result = metric.aggregate() + np.testing.assert_allclose(expected_value, result, rtol=1e-5) + result = metric.aggregate(average=average) # test optional argument metric.reset() np.testing.assert_allclose(expected_value, result, rtol=1e-5) diff --git a/tests/test_config_item.py b/tests/test_config_item.py index fbd76e7be7..4d9df0a870 100644 --- a/tests/test_config_item.py +++ b/tests/test_config_item.py @@ -43,8 +43,10 @@ ] # test execute some function in args, test pre-imported global packages `monai` TEST_CASE_9 = ["collate_fn", "$monai.data.list_data_collate"] -# test lambda function, should not execute the lambda function, just change the string +# test lambda function TEST_CASE_10 = ["collate_fn", "$lambda x: monai.data.list_data_collate(x) + torch.tensor(var)"] +# test regular expression with reference +TEST_CASE_11 = ["collate_fn", "$var + 100"] class TestConfigItem(unittest.TestCase): @@ -72,12 +74,17 @@ def test_component(self, test_input, output_type): if isinstance(ret, LoadImaged): self.assertEqual(ret.keys[0], "image") - @parameterized.expand([TEST_CASE_9, TEST_CASE_10]) + @parameterized.expand([TEST_CASE_9, TEST_CASE_10, TEST_CASE_11]) def test_expression(self, id, test_input): configer = ConfigExpression(id=id, config=test_input, globals={"monai": monai, "torch": torch}) var = 100 - ret = configer.evaluate(locals={"var": var}) - self.assertTrue(isinstance(ret, Callable)) + ret = configer.evaluate(globals={"var": var, "monai": monai}) # `{"monai": monai}` to verify the warning + if isinstance(ret, Callable): + self.assertTrue(isinstance(ret([torch.tensor(1), torch.tensor(2)]), torch.Tensor)) + else: + # also test the `locals` for regular expressions + ret = configer.evaluate(locals={"var": var}) + self.assertEqual(ret, 200) def test_lazy_instantiation(self): config = {"_target_": "DataLoader", "dataset": Dataset(data=[1, 2]), "batch_size": 2} diff --git a/tests/test_config_parser.py b/tests/test_config_parser.py index 9c727c29ac..81e96d4095 100644 --- a/tests/test_config_parser.py +++ b/tests/test_config_parser.py @@ -14,9 +14,11 @@ import unittest from unittest import skipUnless +import numpy as np from parameterized import parameterized -from monai.bundle.config_parser import ConfigParser +from monai.bundle import ConfigParser, ReferenceResolver +from monai.bundle.config_item import ConfigItem from monai.data import DataLoader, Dataset from monai.transforms import Compose, LoadImaged, RandTorchVisiond from monai.utils import min_version, optional_import @@ -86,6 +88,8 @@ def __call__(self, a, b): } ] +TEST_CASE_4 = [{"A": 1, "B": "@A", "C": "@D", "E": "$'test' + '@F'"}] + class TestConfigParser(unittest.TestCase): def test_config_content(self): @@ -111,7 +115,12 @@ def test_parse(self, config, expected_ids, output_types): parser = ConfigParser(config=config, globals={"monai": "monai"}) # test lazy instantiation with original config content parser["transform"]["transforms"][0]["keys"] = "label1" - self.assertEqual(parser.get_parsed_content(id="transform#transforms#0").keys[0], "label1") + trans = parser.get_parsed_content(id="transform#transforms#0") + self.assertEqual(trans.keys[0], "label1") + # test re-use the parsed content or not with the `lazy` option + self.assertEqual(trans, parser.get_parsed_content(id="transform#transforms#0")) + self.assertEqual(trans, parser.get_parsed_content(id="transform#transforms#0", lazy=True)) + self.assertNotEqual(trans, parser.get_parsed_content(id="transform#transforms#0", lazy=False)) # test nested id parser["transform#transforms#0#keys"] = "label2" self.assertEqual(parser.get_parsed_content(id="transform#transforms#0").keys[0], "label2") @@ -121,6 +130,8 @@ def test_parse(self, config, expected_ids, output_types): root = parser.get_parsed_content(id="") for v, cls in zip(root.values(), [Compose, Dataset, DataLoader]): self.assertTrue(isinstance(v, cls)) + # test default value + self.assertEqual(parser.get_parsed_content(id="abc", default=ConfigItem(12345, "abc")), 12345) @parameterized.expand([TEST_CASE_2]) def test_function(self, config): @@ -154,6 +165,71 @@ def test_macro_replace(self): parser.resolve_macro_and_relative_ids() self.assertEqual(str(parser.get()), str({"A": {"B": 1, "C": 2}, "D": [3, 1, 3, 4]})) + @parameterized.expand([TEST_CASE_4]) + def test_allow_missing_reference(self, config): + default = ReferenceResolver.allow_missing_reference + ReferenceResolver.allow_missing_reference = True + parser = ConfigParser(config=config) + + for id in config: + item = parser.get_parsed_content(id=id) + if id in ("A", "B"): + self.assertEqual(item, 1) + elif id == "C": + self.assertEqual(item, "@D") + elif id == "E": + self.assertEqual(item, "test@F") + + # restore the default value + ReferenceResolver.allow_missing_reference = default + with self.assertRaises(ValueError): + parser.parse() + parser.get_parsed_content(id="E") + + def test_list_expressions(self): + config = { + "transform": { + "_target_": "Compose", + "transforms": [{"_target_": "RandScaleIntensity", "factors": 0.5, "prob": 1.0}], + }, + "training": ["$monai.utils.set_determinism(seed=123)", "$@transform(np.asarray([1, 2]))"], + } + parser = ConfigParser(config=config) + parser.get_parsed_content("training", lazy=True, instantiate=True, eval_expr=True) + np.testing.assert_allclose(parser.get_parsed_content("training#1", lazy=True), [0.7942, 1.5885], atol=1e-4) + + def test_contains(self): + empty_parser = ConfigParser({}) + empty_parser.parse() + + parser = ConfigParser({"value": 1, "entry": "string content"}) + parser.parse() + + with self.subTest("Testing empty parser"): + self.assertFalse("something" in empty_parser) + + with self.subTest("Testing with keys"): + self.assertTrue("value" in parser) + self.assertFalse("value1" in parser) + self.assertTrue("entry" in parser) + self.assertFalse("entr" in parser) + + def test_lambda_reference(self): + configs = { + "patch_size": [8, 8], + "transform": {"_target_": "Lambda", "func": "$lambda x: x.reshape((1, *@patch_size))"}, + } + parser = ConfigParser(config=configs) + trans = parser.get_parsed_content(id="transform") + result = trans(np.ones(64)) + self.assertTupleEqual(result.shape, (1, 8, 8)) + + def test_error_instance(self): + config = {"transform": {"_target_": "Compose", "transforms_wrong_key": []}} + parser = ConfigParser(config=config) + with self.assertRaises(RuntimeError): + parser.get_parsed_content("transform", instantiate=True, eval_expr=True) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_convert_data_type.py b/tests/test_convert_data_type.py index 1818f500f9..ab4bd3e3e6 100644 --- a/tests/test_convert_data_type.py +++ b/tests/test_convert_data_type.py @@ -16,17 +16,18 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor from monai.utils.type_conversion import convert_data_type, convert_to_dst_type -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL TESTS: List[Tuple] = [] -for in_type in TEST_NDARRAYS + (int, float): - for out_type in TEST_NDARRAYS: +for in_type in TEST_NDARRAYS_ALL + (int, float): + for out_type in TEST_NDARRAYS_ALL: TESTS.append((in_type(np.array(1.0)), out_type(np.array(1.0)))) # type: ignore TESTS_LIST: List[Tuple] = [] -for in_type in TEST_NDARRAYS + (int, float): - for out_type in TEST_NDARRAYS: +for in_type in TEST_NDARRAYS_ALL + (int, float): + for out_type in TEST_NDARRAYS_ALL: TESTS_LIST.append( ([in_type(np.array(1.0)), in_type(np.array(1.0))], out_type(np.array([1.0, 1.0])), True) # type: ignore ) @@ -83,14 +84,17 @@ def test_convert_data_type(self, in_image, im_out): self.assertEqual(type(in_image), orig_type) if isinstance(in_image, torch.Tensor): self.assertEqual(in_image.device, orig_device) + # check output is desired type - if isinstance(im_out, torch.Tensor): + if isinstance(im_out, MetaTensor): + output_type = MetaTensor + elif isinstance(im_out, torch.Tensor): output_type = torch.Tensor else: output_type = np.ndarray self.assertEqual(type(converted_im), output_type) # check dtype is unchanged - if isinstance(in_type, (np.ndarray, torch.Tensor)): + if isinstance(in_type, (np.ndarray, torch.Tensor, MetaTensor)): self.assertEqual(converted_im.dtype, im_out.dtype) diff --git a/tests/test_crop_foreground.py b/tests/test_crop_foreground.py index af945673fe..e400406e4d 100644 --- a/tests/test_crop_foreground.py +++ b/tests/test_crop_foreground.py @@ -12,15 +12,15 @@ import unittest import numpy as np -import torch from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms import CropForeground -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TEST_COORDS, TESTS = [], [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_COORDS.append( [ {"select_fn": lambda x: x > 0, "channel_indices": None, "margin": 0}, @@ -89,8 +89,15 @@ class TestCropForeground(unittest.TestCase): @parameterized.expand(TEST_COORDS + TESTS) def test_value(self, argments, image, expected_data): - result = CropForeground(**argments)(image) - torch.testing.assert_allclose(result, expected_data, rtol=1e-7, atol=0) + cropper = CropForeground(**argments) + result = cropper(image) + assert_allclose(result, expected_data, type_test=False) + self.assertIsInstance(result, MetaTensor) + self.assertEqual(len(result.applied_operations), 1) + inv = cropper.inverse(result) + self.assertIsInstance(inv, MetaTensor) + self.assertEqual(inv.applied_operations, []) + self.assertTupleEqual(inv.shape, image.shape) @parameterized.expand(TEST_COORDS) def test_return_coords(self, argments, image, _): diff --git a/tests/test_crop_foregroundd.py b/tests/test_crop_foregroundd.py index fa69143827..ab42d6694d 100644 --- a/tests/test_crop_foregroundd.py +++ b/tests/test_crop_foregroundd.py @@ -12,14 +12,13 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import CropForegroundd -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TEST_POSITION, TESTS = [], [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_POSITION.append( [ @@ -151,12 +150,15 @@ class TestCropForegroundd(unittest.TestCase): @parameterized.expand(TEST_POSITION + TESTS) def test_value(self, argments, input_data, expected_data): - result = CropForegroundd(**argments)(input_data) - r, i = result["img"], input_data["img"] - self.assertEqual(type(r), type(i)) - if isinstance(r, torch.Tensor): - self.assertEqual(r.device, i.device) - assert_allclose(r, expected_data) + cropper = CropForegroundd(**argments) + result = cropper(input_data) + assert_allclose(result["img"], expected_data, type_test="tensor") + if "label" in input_data and "img" in input_data: + self.assertTupleEqual(result["img"].shape, result["label"].shape) + inv = cropper.inverse(result) + self.assertTupleEqual(inv["img"].shape, input_data["img"].shape) + if "label" in input_data: + self.assertTupleEqual(inv["label"].shape, input_data["label"].shape) @parameterized.expand(TEST_POSITION) def test_foreground_position(self, argments, input_data, _): diff --git a/tests/test_cross_validation.py b/tests/test_cross_validation.py index c378a52f78..811dcea026 100644 --- a/tests/test_cross_validation.py +++ b/tests/test_cross_validation.py @@ -13,8 +13,8 @@ import unittest from monai.apps import CrossValidation, DecathlonDataset -from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd, ToTensord -from monai.utils.enums import PostFix +from monai.data import MetaTensor +from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd from tests.utils import skip_if_downloading_fails, skip_if_quick @@ -23,12 +23,7 @@ class TestCrossValidation(unittest.TestCase): def test_values(self): testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") train_transform = Compose( - [ - LoadImaged(keys=["image", "label"]), - AddChanneld(keys=["image", "label"]), - ScaleIntensityd(keys="image"), - ToTensord(keys=["image", "label"]), - ] + [LoadImaged(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), ScaleIntensityd(keys="image")] ) val_transform = LoadImaged(keys=["image", "label"]) @@ -36,7 +31,7 @@ def _test_dataset(dataset): self.assertEqual(len(dataset), 52) self.assertTrue("image" in dataset[0]) self.assertTrue("label" in dataset[0]) - self.assertTrue(PostFix.meta("image") in dataset[0]) + self.assertTrue(isinstance(dataset[0]["image"], MetaTensor)) self.assertTupleEqual(dataset[0]["image"].shape, (1, 34, 49, 41)) cvdataset = CrossValidation( diff --git a/tests/test_cumulative.py b/tests/test_cumulative.py index 12a6a5e5e7..16f5c1d1f5 100644 --- a/tests/test_cumulative.py +++ b/tests/test_cumulative.py @@ -35,6 +35,8 @@ def test_multi(self): c.append() c.extend() self.assertEqual(c.get_buffer(), []) + c.get_buffer().append(1) + self.assertEqual(c.get_buffer(), []) # no in-place update for the buffer c.reset() diff --git a/tests/test_data_stats.py b/tests/test_data_stats.py index c18abfcedc..2b652f8f62 100644 --- a/tests/test_data_stats.py +++ b/tests/test_data_stats.py @@ -14,6 +14,8 @@ import sys import tempfile import unittest +from io import StringIO +from unittest.mock import patch import numpy as np import torch @@ -137,7 +139,6 @@ class TestDataStats(unittest.TestCase): def test_value(self, input_param, input_data, expected_print): transform = DataStats(**input_param) _ = transform(input_data) - # self.assertEqual(transform.output, expected_print) @parameterized.expand([TEST_CASE_8]) def test_file(self, input_data, expected_print): @@ -167,6 +168,14 @@ def test_file(self, input_data, expected_print): if sys.platform != "win32": self.assertEqual(content, expected_print) + def test_multiple_data_stats(self): + with patch("sys.stdout", new=StringIO()) as out: + input_data = np.array([[0, 1], [1, 2]]) + transform = DataStats() + _ = DataStats() + _ = transform(input_data) + print(out.getvalue().strip()) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_data_statsd.py b/tests/test_data_statsd.py index 28da936cd0..9c878addf5 100644 --- a/tests/test_data_statsd.py +++ b/tests/test_data_statsd.py @@ -161,7 +161,6 @@ class TestDataStatsd(unittest.TestCase): def test_value(self, input_param, input_data, expected_print): transform = DataStatsd(**input_param) _ = transform(input_data) - # self.assertEqual(transform.printer.output, expected_print) @parameterized.expand([TEST_CASE_9]) def test_file(self, input_data, expected_print): diff --git a/tests/test_dataset_summary.py b/tests/test_dataset_summary.py index 51840f77ea..d0531b28a0 100644 --- a/tests/test_dataset_summary.py +++ b/tests/test_dataset_summary.py @@ -19,6 +19,9 @@ from monai.data import Dataset, DatasetSummary, create_test_image_3d from monai.transforms import LoadImaged +from monai.transforms.compose import Compose +from monai.transforms.meta_utility.dictionary import FromMetaTensord +from monai.transforms.utility.dictionary import ToNumpyd from monai.utils import set_determinism from monai.utils.enums import PostFix @@ -50,12 +53,17 @@ def test_spacing_intensity(self): {"image": image_name, "label": label_name} for image_name, label_name in zip(train_images, train_labels) ] - dataset = Dataset( - data=data_dicts, transform=LoadImaged(keys=["image", "label"], meta_keys=["test1", "test2"]) + t = Compose( + [ + LoadImaged(keys=["image", "label"]), + FromMetaTensord(keys=["image", "label"]), + ToNumpyd(keys=["image", "label", "image_meta_dict", "label_meta_dict"]), + ] ) + dataset = Dataset(data=data_dicts, transform=t) # test **kwargs of `DatasetSummary` for `DataLoader` - calculator = DatasetSummary(dataset, num_workers=4, meta_key="test1", collate_fn=test_collate) + calculator = DatasetSummary(dataset, num_workers=4, meta_key="image_meta_dict", collate_fn=test_collate) target_spacing = calculator.get_target_spacing() self.assertEqual(target_spacing, (1.0, 1.0, 1.0)) @@ -85,7 +93,8 @@ def test_anisotropic_spacing(self): {"image": image_name, "label": label_name} for image_name, label_name in zip(train_images, train_labels) ] - dataset = Dataset(data=data_dicts, transform=LoadImaged(keys=["image", "label"])) + t = Compose([LoadImaged(keys=["image", "label"]), FromMetaTensord(keys=["image", "label"])]) + dataset = Dataset(data=data_dicts, transform=t) calculator = DatasetSummary(dataset, num_workers=4, meta_key_postfix=PostFix.meta()) diff --git a/tests/test_decathlondataset.py b/tests/test_decathlondataset.py index 744dccefaa..49280f6fa6 100644 --- a/tests/test_decathlondataset.py +++ b/tests/test_decathlondataset.py @@ -15,8 +15,8 @@ from pathlib import Path from monai.apps import DecathlonDataset -from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd, ToTensord -from monai.utils.enums import PostFix +from monai.data import MetaTensor +from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd from tests.utils import skip_if_downloading_fails, skip_if_quick @@ -25,19 +25,14 @@ class TestDecathlonDataset(unittest.TestCase): def test_values(self): testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") transform = Compose( - [ - LoadImaged(keys=["image", "label"]), - AddChanneld(keys=["image", "label"]), - ScaleIntensityd(keys="image"), - ToTensord(keys=["image", "label"]), - ] + [LoadImaged(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), ScaleIntensityd(keys="image")] ) def _test_dataset(dataset): self.assertEqual(len(dataset), 52) self.assertTrue("image" in dataset[0]) self.assertTrue("label" in dataset[0]) - self.assertTrue(PostFix.meta("image") in dataset[0]) + self.assertTrue(isinstance(dataset[0]["image"], MetaTensor)) self.assertTupleEqual(dataset[0]["image"].shape, (1, 36, 47, 44)) with skip_if_downloading_fails(): @@ -55,8 +50,8 @@ def _test_dataset(dataset): root_dir=testing_dir, task="Task04_Hippocampus", transform=transform, section="validation", download=False ) _test_dataset(data) - self.assertTrue(data[0][PostFix.meta("image")]["filename_or_obj"].endswith("hippocampus_163.nii.gz")) - self.assertTrue(data[0][PostFix.meta("label")]["filename_or_obj"].endswith("hippocampus_163.nii.gz")) + self.assertTrue(data[0]["image"].meta["filename_or_obj"].endswith("hippocampus_163.nii.gz")) + self.assertTrue(data[0]["label"].meta["filename_or_obj"].endswith("hippocampus_163.nii.gz")) # test validation without transforms data = DecathlonDataset(root_dir=testing_dir, task="Task04_Hippocampus", section="validation", download=False) self.assertTupleEqual(data[0]["image"].shape, (36, 47, 44)) diff --git a/tests/test_decollate.py b/tests/test_decollate.py index adeaa73337..440f5e59e7 100644 --- a/tests/test_decollate.py +++ b/tests/test_decollate.py @@ -74,19 +74,12 @@ ], [[None, None], [None, None]], [["test"], ["test"]], + [np.array([64, 64]), [64, 64]], [[], []], [[("ch1", "ch2"), ("ch3",)], [["ch1", "ch3"], ["ch2", None]]], # default pad None ] -class _ListCompose(Compose): - def __call__(self, input_): - img, metadata = self.transforms[0](input_) - for t in self.transforms[1:]: - img = t(img) - return img, metadata - - class TestDeCollate(unittest.TestCase): def setUp(self) -> None: set_determinism(seed=0) @@ -138,7 +131,7 @@ def test_decollation_dict(self, *transforms): t_compose = Compose([AddChanneld(KEYS), Compose(transforms), ToTensord(KEYS)]) # If nibabel present, read from disk if has_nib: - t_compose = Compose([LoadImaged("image"), t_compose]) + t_compose = Compose([LoadImaged("image", image_only=True), t_compose]) dataset = CacheDataset(self.data_dict, t_compose, progress=False) self.check_decollate(dataset=dataset) @@ -158,7 +151,7 @@ def test_decollation_list(self, *transforms): t_compose = Compose([AddChannel(), Compose(transforms), ToTensor()]) # If nibabel present, read from disk if has_nib: - t_compose = _ListCompose([LoadImage(image_only=False), t_compose]) + t_compose = Compose([LoadImage(image_only=True), t_compose]) dataset = Dataset(self.data_list, t_compose) self.check_decollate(dataset=dataset) diff --git a/tests/test_deepedit_interaction.py b/tests/test_deepedit_interaction.py new file mode 100644 index 0000000000..6bb723268f --- /dev/null +++ b/tests/test_deepedit_interaction.py @@ -0,0 +1,119 @@ +# 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. + +import unittest + +import numpy as np +import torch + +from monai.apps.deepedit.interaction import Interaction +from monai.apps.deepedit.transforms import ( + AddGuidanceSignalDeepEditd, + AddInitialSeedPointMissingLabelsd, + AddRandomGuidanceDeepEditd, + FindAllValidSlicesMissingLabelsd, + FindDiscrepancyRegionsDeepEditd, + SplitPredsLabeld, +) +from monai.data import Dataset +from monai.engines import SupervisedTrainer +from monai.engines.utils import IterationEvents +from monai.losses import DiceCELoss +from monai.transforms import Activationsd, AsDiscreted, Compose, ToTensord + + +def add_one(engine): + if engine.state.best_metric == -1: + engine.state.best_metric = 0 + else: + engine.state.best_metric = engine.state.best_metric + 1 + + +class TestInteractions(unittest.TestCase): + def run_interaction(self, train): + label_names = {"spleen": 1, "background": 0} + np.random.seed(0) + data = [ + { + "image": np.random.randint(0, 256, size=(1, 15, 15, 15)).astype(np.float32), + "label": np.random.randint(0, 2, size=(1, 15, 15, 15)), + "label_names": label_names, + } + for _ in range(5) + ] + network = torch.nn.Conv3d(3, len(label_names), 1) + lr = 1e-3 + opt = torch.optim.Adam(network.parameters(), lr) + loss = DiceCELoss(to_onehot_y=True, softmax=True) + pre_transforms = Compose( + [ + FindAllValidSlicesMissingLabelsd(keys="label", sids="sids"), + AddInitialSeedPointMissingLabelsd(keys="label", guidance="guidance", sids="sids"), + AddGuidanceSignalDeepEditd(keys="image", guidance="guidance", number_intensity_ch=1), + ToTensord(keys=("image", "label")), + ] + ) + dataset = Dataset(data, transform=pre_transforms) + data_loader = torch.utils.data.DataLoader(dataset, batch_size=5) + + iteration_transforms = [ + FindDiscrepancyRegionsDeepEditd(keys="label", pred="pred", discrepancy="discrepancy"), + AddRandomGuidanceDeepEditd( + keys="NA", guidance="guidance", discrepancy="discrepancy", probability="probability" + ), + AddGuidanceSignalDeepEditd(keys="image", guidance="guidance", number_intensity_ch=1), + ToTensord(keys=("image", "label")), + ] + post_transforms = [ + Activationsd(keys="pred", softmax=True), + AsDiscreted(keys=("pred", "label"), argmax=(True, False), to_onehot=len(label_names)), + SplitPredsLabeld(keys="pred"), + ToTensord(keys=("image", "label")), + ] + iteration_transforms = Compose(iteration_transforms) + post_transforms = Compose(post_transforms) + + i = Interaction( + deepgrow_probability=1.0, + transforms=iteration_transforms, + click_probability_key="probability", + train=train, + label_names=label_names, + ) + self.assertEqual(len(i.transforms.transforms), 4, "Mismatch in expected transforms") + + # set up engine + engine = SupervisedTrainer( + device=torch.device("cpu"), + max_epochs=1, + train_data_loader=data_loader, + network=network, + optimizer=opt, + loss_function=loss, + postprocessing=post_transforms, + iteration_update=i, + ) + engine.add_event_handler(IterationEvents.INNER_ITERATION_STARTED, add_one) + engine.add_event_handler(IterationEvents.INNER_ITERATION_COMPLETED, add_one) + + engine.run() + self.assertIsNotNone(engine.state.batch[0].get("guidance"), "guidance is missing") + self.assertEqual(engine.state.best_metric, 1) + + def test_train_interaction(self): + self.run_interaction(train=True) + + def test_val_interaction(self): + self.run_interaction(train=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_deepedit_transforms.py b/tests/test_deepedit_transforms.py index a5c5f0fe2f..225b2fc60b 100644 --- a/tests/test_deepedit_transforms.py +++ b/tests/test_deepedit_transforms.py @@ -14,81 +14,288 @@ import numpy as np from parameterized import parameterized -from monai.apps.deepedit.transforms import ClickRatioAddRandomGuidanced, DiscardAddGuidanced, ResizeGuidanceCustomd +from monai.apps.deepedit.transforms import ( + AddGuidanceFromPointsDeepEditd, + AddGuidanceSignalDeepEditd, + AddInitialSeedPointMissingLabelsd, + AddRandomGuidanceDeepEditd, + DiscardAddGuidanced, + FindAllValidSlicesMissingLabelsd, + FindDiscrepancyRegionsDeepEditd, + NormalizeLabelsInDatasetd, + ResizeGuidanceMultipleLabelDeepEditd, + SingleLabelSelectiond, + SplitPredsLabeld, +) +from monai.utils import min_version, optional_import, set_determinism from monai.utils.enums import PostFix -IMAGE = np.array([[[[1, 0, 2, 0, 1], [0, 1, 2, 1, 0], [2, 2, 3, 2, 2], [0, 1, 2, 1, 0], [1, 0, 2, 0, 1]]]]) -LABEL = np.array([[[[0, 0, 0, 0, 0], [0, 1, 0, 1, 0], [0, 0, 1, 0, 0], [0, 1, 0, 1, 0], [0, 0, 0, 0, 0]]]]) +measure, _ = optional_import("skimage.measure", "0.14.2", min_version) + +set_determinism(seed=0) +IMAGE = np.random.randint(0, 256, size=(1, 10, 10, 10)) +THREE_CHAN_IMAGE = np.random.randint(0, 255, size=(3, 10, 10, 10)) +LABEL = np.random.randint(0, 2, size=(10, 10, 10)) +PRED = np.random.randint(0, 2, size=(10, 10, 10)) +LABEL_NAMES = {"spleen": 1, "background": 0} +DISCREPANCY = { + "spleen": np.random.randint(0, 2, size=(10, 10, 10)), + "background": np.random.randint(0, 2, size=(10, 10, 10)), +} +set_determinism(None) DATA_1 = { "image": IMAGE, "label": LABEL, - PostFix.meta("image"): {"dim": IMAGE.shape}, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, PostFix.meta("label"): {}, - "foreground": [0, 0, 0], - "background": [0, 0, 0], } -DISCARD_ADD_GUIDANCE_TEST_CASE = [{"image": IMAGE, "label": LABEL}, DATA_1, (3, 1, 5, 5)] - DATA_2 = { "image": IMAGE, "label": LABEL, - "guidance": np.array([[[1, 0, 2, 2]], [[-1, -1, -1, -1]]]), - "discrepancy": np.array( - [ - [[[[0, 0, 0, 0, 0], [0, 0, 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, 1, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]]], - ] - ), + "label_names": LABEL_NAMES, + "guidance": {"spleen": [[3, 5, 4, 6], [-1, -1, -1, -1]], "background": [[-1, -1, -1, -1], [-1, -1, -1, -1]]}, + "discrepancy": DISCREPANCY, "probability": 1.0, } -CLICK_RATIO_ADD_RANDOM_GUIDANCE_TEST_CASE_1 = [ - {"guidance": "guidance", "discrepancy": "discrepancy", "probability": "probability"}, - DATA_2, - "[[[1, 0, 2, 2], [-1, -1, -1, -1]], [[-1, -1, -1, -1], [1, 0, 2, 1]]]", -] - DATA_3 = { - "image": np.arange(1000).reshape((1, 5, 10, 20)), - PostFix.meta("image"): {"foreground_cropped_shape": (1, 10, 20, 40), "dim": [3, 512, 512, 128]}, - "guidance": [[[6, 10, 14], [8, 10, 14]], [[8, 10, 16]]], - "foreground": [[10, 14, 6], [10, 14, 8]], - "background": [[10, 16, 8]], + "image": IMAGE, + "label": LABEL, + "guidance": { + "spleen": np.array([[1, 0, 2, 2], [-1, -1, -1, -1]]), + "background": np.array([[1, 0, 2, 2], [-1, -1, -1, -1]]), + }, + "probability": 1.0, + "label_names": LABEL_NAMES, + "pred": PRED, +} + +DATA_4 = { + "image": IMAGE, + "label": LABEL, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "current_label": "spleen", + "probability": 1.0, + "label_names": LABEL_NAMES, + "spleen": [[0, 4, 3], [0, 0, 3], [0, 1, 3]], + "sids": {"spleen": []}, + "pred": PRED, +} + +DATA_5 = { + "image": IMAGE, + "label": LABEL, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "current_label": "spleen", + "probability": 1.0, + "label_names": LABEL_NAMES, + "sids": {"spleen": [2, 3, 4], "background": [0, 1, 5]}, +} + +DATA_6 = { + "image": IMAGE, + "label": LABEL[None], + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "current_label": "spleen", + "label_names": LABEL_NAMES, +} + +DATA_7 = { + "image": IMAGE, + "label": LABEL, + "pred": PRED, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "current_label": "spleen", + "probability": 1.0, + "label_names": LABEL_NAMES, + "guidance": { + "spleen": np.array([[1, 0, 2, 2], [-1, -1, -1, -1]]), + "background": np.array([[1, 0, 2, 2], [-1, -1, -1, -1]]), + }, +} + +DATA_8 = { + "image": IMAGE, + "label": LABEL, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "label_names": LABEL_NAMES, +} + +DATA_9 = { + "image": IMAGE, + "label": LABEL, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "label_names": LABEL_NAMES, + "guidance": {"spleen": np.array([0, 2, 2]), "background": np.array([-1, -1, -1])}, +} + +DATA_10 = { + "image": IMAGE, + "label": LABEL, + PostFix.meta("image"): {"dim": IMAGE.shape, "spatial_shape": IMAGE[0, ...].shape}, + "current_label": "spleen", } -RESIZE_GUIDANCE_TEST_CASE_1 = [ - {"ref_image": "image", "guidance": "guidance"}, - DATA_3, - [[[0, 0, 0], [0, 0, 1]], [[0, 0, 1]]], +DATA_11 = {"image": IMAGE, "label": LABEL, "label_names": LABEL_NAMES, "pred": PRED} + + +ADD_GUIDANCE_FROM_POINTS_TEST_CASE = [ + {"ref_image": "image", "guidance": "guidance", "label_names": LABEL_NAMES}, # arguments + DATA_4, # input_data + [0, 4, 3], # expected_result ] +ADD_GUIDANCE_CUSTOM_TEST_CASE = [ + {"keys": "image", "guidance": "guidance"}, # arguments + DATA_3, # input_data + 3, # expected_result +] -class TestDiscardAddGuidanced(unittest.TestCase): - @parameterized.expand([DISCARD_ADD_GUIDANCE_TEST_CASE]) +ADD_INITIAL_POINT_TEST_CASE = [ + {"keys": "label", "guidance": "guidance", "sids": "sids"}, # arguments + DATA_5, # input_data + { + "spleen": "[[1, 0, 7], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1]]", + "background": "[[1, 5, 3], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1]]", + }, # expected_result +] + +ADD_RANDOM_GUIDANCE_TEST_CASE = [ + {"keys": "NA", "guidance": "guidance", "discrepancy": "discrepancy", "probability": "probability"}, # arguments + DATA_2, # input_data + {"spleen": [[3, 5, 4, 6], [-1, -1, -1, -1]], "background": [[-1, -1, -1, -1], [-1, -1, -1, -1]]}, # expected_result +] + +DISCARD_ADD_GUIDANCE_TEST_CASE = [ + {"keys": "image", "label_names": LABEL_NAMES}, # arguments + DATA_1, # input_data + (3, 10, 10, 10), # expected_result +] + +FIND_DISCREPANCY_TEST_CASE = [ + {"keys": "label", "pred": "pred", "discrepancy": "discrepancy"}, # arguments + DATA_7, # input_data + 240, # expected_result +] + +FIND_SLICE_TEST_CASE = [ + {"keys": "label", "sids": "sids"}, # arguments + DATA_6, # input_data + {"spleen": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], "background": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}, # expected_result +] + +NormalizeLabelsDatasetd_TEST_CASE = [ + {"keys": "label", "label_names": LABEL_NAMES}, # arguments + DATA_8, # input_data + len(LABEL_NAMES), # expected_result +] + +RESIZE_GUIDANCE_TEST_CASE = [ + {"guidance": "guidance", "ref_image": "image"}, # arguments + DATA_9, # input_data + {"spleen": [0, 2, 2], "background": [-1, -1, -1]}, # expected_result +] + +SingleLabelSelectiond_TEST_CASE = [ + {"keys": "label", "label_names": ["spleen"]}, # arguments + DATA_10, # input_data + "spleen", # expected_result +] + +SplitPredsLabeld_TEST_CASE = [{"keys": "pred"}, DATA_11, (1, 10, 10)] # arguments # input_data # expected_result + + +class TestAddGuidanceFromPointsCustomd(unittest.TestCase): + @parameterized.expand([ADD_GUIDANCE_FROM_POINTS_TEST_CASE]) def test_correct_results(self, arguments, input_data, expected_result): - add_fn = DiscardAddGuidanced(arguments) + add_fn = AddGuidanceFromPointsDeepEditd(**arguments) result = add_fn(input_data) - self.assertEqual(result["image"].shape, expected_result) + self.assertEqual(result[arguments["guidance"]]["spleen"][0], expected_result) -class TestClickRatioAddRandomGuidanced(unittest.TestCase): - @parameterized.expand([CLICK_RATIO_ADD_RANDOM_GUIDANCE_TEST_CASE_1]) +class TestAddGuidanceSignalCustomd(unittest.TestCase): + @parameterized.expand([ADD_GUIDANCE_CUSTOM_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = AddGuidanceSignalDeepEditd(**arguments) + result = add_fn(input_data) + self.assertEqual(result["image"].shape[0], expected_result) + + +class TestAddInitialSeedPointMissingLabelsd(unittest.TestCase): + @parameterized.expand([ADD_INITIAL_POINT_TEST_CASE]) def test_correct_results(self, arguments, input_data, expected_result): seed = 0 - add_fn = ClickRatioAddRandomGuidanced(**arguments) + add_fn = AddInitialSeedPointMissingLabelsd(**arguments) add_fn.set_random_state(seed) result = add_fn(input_data) self.assertEqual(result[arguments["guidance"]], expected_result) -class TestResizeGuidanced(unittest.TestCase): - @parameterized.expand([RESIZE_GUIDANCE_TEST_CASE_1]) +class TestAddRandomGuidanceCustomd(unittest.TestCase): + @parameterized.expand([ADD_RANDOM_GUIDANCE_TEST_CASE]) def test_correct_results(self, arguments, input_data, expected_result): - result = ResizeGuidanceCustomd(**arguments)(input_data) + add_fn = AddRandomGuidanceDeepEditd(**arguments) + result = add_fn(input_data) self.assertEqual(result[arguments["guidance"]], expected_result) +class TestDiscardAddGuidanced(unittest.TestCase): + @parameterized.expand([DISCARD_ADD_GUIDANCE_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = DiscardAddGuidanced(**arguments) + result = add_fn(input_data) + self.assertEqual(result["image"].shape, expected_result) + + +class TestFindAllValidSlicesMissingLabelsd(unittest.TestCase): + @parameterized.expand([FIND_SLICE_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = FindAllValidSlicesMissingLabelsd(**arguments) + result = add_fn(input_data) + self.assertEqual(result[arguments["sids"]], expected_result) + + +class TestFindDiscrepancyRegionsCustomd(unittest.TestCase): + @parameterized.expand([FIND_DISCREPANCY_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = FindDiscrepancyRegionsDeepEditd(**arguments) + result = add_fn(input_data) + self.assertEqual(np.sum(result[arguments["discrepancy"]]["spleen"][0]), expected_result) + + +class TestNormalizeLabelsDatasetd(unittest.TestCase): + @parameterized.expand([NormalizeLabelsDatasetd_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = NormalizeLabelsInDatasetd(**arguments) + result = add_fn(input_data) + self.assertEqual(len(np.unique(result["label"])), expected_result) + + +class TestResizeGuidanceMultipleLabelCustomd(unittest.TestCase): + @parameterized.expand([RESIZE_GUIDANCE_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = ResizeGuidanceMultipleLabelDeepEditd(**arguments) + result = add_fn(input_data) + self.assertEqual(result[arguments["guidance"]], expected_result) + + +class TestSingleLabelSelectiond(unittest.TestCase): + @parameterized.expand([SingleLabelSelectiond_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = SingleLabelSelectiond(**arguments) + result = add_fn(input_data) + self.assertEqual(result["current_label"], expected_result) + + +class TestSplitPredsLabeld(unittest.TestCase): + @parameterized.expand([SplitPredsLabeld_TEST_CASE]) + def test_correct_results(self, arguments, input_data, expected_result): + add_fn = SplitPredsLabeld(**arguments) + result = add_fn(input_data) + self.assertEqual(result["pred_spleen"].shape, expected_result) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_delete_itemsd.py b/tests/test_delete_itemsd.py index 450abb40db..99d05fe787 100644 --- a/tests/test_delete_itemsd.py +++ b/tests/test_delete_itemsd.py @@ -48,7 +48,7 @@ def test_re(self, input_param): input_data = {"image": [1, 2, 3], PostFix.meta(): {"0008|0005": 1, "0008|1050": 2, "0008test": 3}} result = DeleteItemsd(**input_param)(input_data) self.assertEqual(result[PostFix.meta()]["0008test"], 3) - self.assertTrue(len(result[PostFix.meta()]), 1) + self.assertEqual(len(result[PostFix.meta()]), 1) if __name__ == "__main__": diff --git a/tests/test_detect_envelope.py b/tests/test_detect_envelope.py index f1b9c7ad1a..5eea0c8653 100644 --- a/tests/test_detect_envelope.py +++ b/tests/test_detect_envelope.py @@ -16,8 +16,8 @@ from parameterized import parameterized from monai.transforms import DetectEnvelope -from monai.utils import InvalidPyTorchVersionError, OptionalImportError -from tests.utils import SkipIfAtLeastPyTorchVersion, SkipIfBeforePyTorchVersion, SkipIfModule, SkipIfNoModule +from monai.utils import OptionalImportError +from tests.utils import TEST_NDARRAYS, SkipIfModule, SkipIfNoModule, assert_allclose n_samples = 500 hann_windowed_sine = np.sin(2 * np.pi * 10 * np.linspace(0, 1, n_samples)) * np.hanning(n_samples) @@ -112,7 +112,6 @@ TEST_CASE_INVALID_OBJ = [{}, "a string", "__call__"] # method expected to raise exception -@SkipIfBeforePyTorchVersion((1, 7)) @SkipIfNoModule("torch.fft") class TestDetectEnvelope(unittest.TestCase): @parameterized.expand( @@ -126,8 +125,9 @@ class TestDetectEnvelope(unittest.TestCase): ] ) def test_value(self, arguments, image, expected_data, atol): - result = DetectEnvelope(**arguments)(image) - np.testing.assert_allclose(result, expected_data, atol=atol) + for p in TEST_NDARRAYS: + result = DetectEnvelope(**arguments)(p(image)) + assert_allclose(result, p(expected_data), atol=atol, type_test="tensor") @parameterized.expand( [ @@ -147,19 +147,11 @@ def test_value_error(self, arguments, image, method): raise ValueError("Expected raising method invalid. Should be __init__ or __call__.") -@SkipIfBeforePyTorchVersion((1, 7)) @SkipIfModule("torch.fft") class TestHilbertTransformNoFFTMod(unittest.TestCase): def test_no_fft_module_error(self): self.assertRaises(OptionalImportError, DetectEnvelope(), np.random.rand(1, 10)) -@SkipIfAtLeastPyTorchVersion((1, 7)) -class TestDetectEnvelopeInvalidPyTorch(unittest.TestCase): - def test_invalid_pytorch_error(self): - with self.assertRaisesRegex(InvalidPyTorchVersionError, "version"): - DetectEnvelope() - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_detection_coco_metrics.py b/tests/test_detection_coco_metrics.py new file mode 100644 index 0000000000..b139377511 --- /dev/null +++ b/tests/test_detection_coco_metrics.py @@ -0,0 +1,66 @@ +# 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. + +import random +import unittest + +import numpy as np +import torch + +from monai.apps.detection.metrics.coco import COCOMetric +from monai.apps.detection.metrics.matching import matching_batch +from monai.data.box_utils import box_iou + + +class TestCOCOMetrics(unittest.TestCase): + def test_coco_run(self): + coco_metric = COCOMetric(classes=["c0", "c1", "c2"], iou_list=[0.1], max_detection=[10]) + + num_images = 10 + + val_outputs_all = [] + val_targets_all = [] + for _ in range(num_images): + # randomly generate gt boxes and pred boxes + num_gt_boxes = random.randint(1, 3) + num_pred_boxes = random.randint(0, 3) + + box_start = torch.randint(3, (num_pred_boxes, 3)) + box_stop = box_start + torch.randint(1, 32, (num_pred_boxes, 3)) + boxes = torch.cat((box_start, box_stop), dim=1).to(torch.float16) + val_outputs_all.append( + { + "boxes": boxes, + "labels": torch.randint(3, (num_pred_boxes,)), + "scores": torch.randn((num_pred_boxes,)).absolute(), + } + ) + + box_start = torch.randint(3, (num_gt_boxes, 3)) + box_stop = box_start + torch.randint(1, 32, (num_gt_boxes, 3)) + boxes = torch.cat((box_start, box_stop), dim=1).to(torch.float16) + val_targets_all.append({"boxes": boxes, "labels": torch.randint(3, (num_gt_boxes,))}) + + results_metric = matching_batch( + iou_fn=box_iou, + iou_thresholds=coco_metric.iou_thresholds, + pred_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_outputs_all], + pred_classes=[val_data_i["labels"].numpy() for val_data_i in val_outputs_all], + pred_scores=[val_data_i["scores"].numpy() for val_data_i in val_outputs_all], + gt_boxes=[val_data_i["boxes"].numpy() for val_data_i in val_targets_all], + gt_classes=[val_data_i["labels"].numpy() for val_data_i in val_targets_all], + ) + val_epoch_metric_dict = coco_metric(results_metric)[0] + np.testing.assert_array_less([-16.01], [sum(val_epoch_metric_dict.values())]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_detector_boxselector.py b/tests/test_detector_boxselector.py new file mode 100644 index 0000000000..6e22a7833a --- /dev/null +++ b/tests/test_detector_boxselector.py @@ -0,0 +1,65 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.box_selector import BoxSelector +from tests.utils import assert_allclose + +device = "cuda" if torch.cuda.is_available() else "cpu" +num_anchors = 7 + +TEST_CASE = [] +TEST_CASE.append( + [ # 2D + { + "apply_sigmoid": False, + "score_thresh": 0.1, + "topk_candidates_per_level": 2, + "nms_thresh": 0.1, + "detections_per_img": 5, + }, + [torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 10]]).to(torch.float16)], + [torch.tensor([[0.1, 0.6], [0.2, 0.2]])], + (20, 20, 20), + torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 10]]), + ] +) +TEST_CASE.append( + [ + { + "apply_sigmoid": False, + "score_thresh": 0.1, + "topk_candidates_per_level": 1, + "nms_thresh": 0.1, + "detections_per_img": 5, + }, + [torch.tensor([[1, 2, 3, 2, 3, 4]]).to(torch.float32), torch.tensor([[5, 6, 7, 8, 9, 10]]).to(torch.float32)], + [torch.tensor([[0.3, 0.6]]), torch.tensor([[0.2, 0.2]])], + (20, 20, 8), + torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 8]]), + ] +) + + +class TestBoxSelector(unittest.TestCase): + @parameterized.expand(TEST_CASE) + def test_box_selector(self, input_param, boxes, logits, image_shape, expected_results): + box_selector = BoxSelector(**input_param) + result = box_selector.select_boxes_per_image(boxes, logits, image_shape) + assert_allclose(result[0], expected_results, type_test=True, device_test=False, atol=1e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_detector_utils.py b/tests/test_detector_utils.py new file mode 100644 index 0000000000..b8ae390016 --- /dev/null +++ b/tests/test_detector_utils.py @@ -0,0 +1,94 @@ +# 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. + +import random +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.detector_utils import preprocess_images +from monai.utils import ensure_tuple +from tests.utils import assert_allclose + +TEST_CASE_1 = [ # 3D, batch 3, 2 input channel + { + "pretrained": False, + "spatial_dims": 3, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": 7, + "conv1_t_stride": (2, 2, 2), + }, + (3, 2, 32, 64, 48), + (3, 2, 64, 64, 64), +] + +TEST_CASE_2 = [ # 2D, batch 2, 1 input channel + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "conv1_t_size": [7, 7], + "conv1_t_stride": [2, 2], + }, + (2, 1, 32, 64), + (2, 1, 64, 64), +] + +TEST_CASE_2_A = [ # 2D, batch 2, 1 input channel, shortcut type A + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "shortcut_type": "A", + "conv1_t_size": (7, 7), + "conv1_t_stride": 2, + }, + (2, 1, 32, 64), + (2, 1, 64, 64), +] + +TEST_CASE_3 = [ # 1D, batch 1, 2 input channels + { + "pretrained": False, + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), + (1, 2, 32), +] + +TEST_CASES = [] +TEST_CASES = [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3] + + +class TestDetectorUtils(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_detector_utils(self, input_param, input_shape, expected_shape): + size_divisible = 32 * ensure_tuple(input_param["conv1_t_stride"])[0] + input_data = torch.randn(input_shape) + result, _ = preprocess_images(input_data, input_param["spatial_dims"], size_divisible, mode="constant", value=1) + assert_allclose(expected_shape, result.shape, type_test=True, device_test=False, atol=0.1) + + input_data = [torch.randn(input_shape[1:]) for _ in range(random.randint(1, 9))] + result, _ = preprocess_images(input_data, input_param["spatial_dims"], size_divisible, mode="edge") + expected_shape = (len(input_data),) + expected_shape[1:] + assert_allclose(expected_shape, result.shape, type_test=True, device_test=False, atol=0.1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_dice_ce_loss.py b/tests/test_dice_ce_loss.py index ff2cd00b02..83ad5b8d9a 100644 --- a/tests/test_dice_ce_loss.py +++ b/tests/test_dice_ce_loss.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.losses import DiceCELoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save TEST_CASES = [ [ # shape: (2, 2, 3), (2, 1, 3) @@ -85,7 +85,6 @@ def test_ill_reduction(self): loss = DiceCELoss(reduction="none") loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3))) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = DiceCELoss() test_input = torch.ones(2, 1, 8, 8) diff --git a/tests/test_dice_focal_loss.py b/tests/test_dice_focal_loss.py index c611fe4160..b77a36e720 100644 --- a/tests/test_dice_focal_loss.py +++ b/tests/test_dice_focal_loss.py @@ -15,7 +15,7 @@ import torch from monai.losses import DiceFocalLoss, DiceLoss, FocalLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save class TestDiceFocalLoss(unittest.TestCase): @@ -61,7 +61,6 @@ def test_ill_lambda(self): with self.assertRaisesRegex(ValueError, ""): DiceFocalLoss(lambda_dice=-1.0) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = DiceFocalLoss() test_input = torch.ones(2, 1, 8, 8) diff --git a/tests/test_dice_loss.py b/tests/test_dice_loss.py index 4e45393de6..223b09e624 100644 --- a/tests/test_dice_loss.py +++ b/tests/test_dice_loss.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.losses import DiceLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -184,7 +184,6 @@ def test_input_warnings(self): loss = DiceLoss(to_onehot_y=True) loss.forward(chn_input, chn_target) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = DiceLoss() test_input = torch.ones(2, 1, 8, 8) diff --git a/tests/test_dints_cell.py b/tests/test_dints_cell.py index d480235b70..a5da39bae9 100644 --- a/tests/test_dints_cell.py +++ b/tests/test_dints_cell.py @@ -32,21 +32,28 @@ (2, 4, 64, 32, 16), ], [ - {"c_prev": 8, "c": 8, "rate": 0, "arch_code_c": None}, + {"c_prev": 8, "c": 8, "rate": 0, "arch_code_c": None, "act_name": "SELU", "norm_name": "BATCH"}, torch.tensor([1, 1, 1, 1, 1]), torch.tensor([0, 0, 0, 1, 0]), (2, 8, 32, 16, 8), (2, 8, 32, 16, 8), ], [ - {"c_prev": 8, "c": 8, "rate": -1, "arch_code_c": None}, + { + "c_prev": 8, + "c": 8, + "rate": -1, + "arch_code_c": None, + "act_name": "PRELU", + "norm_name": ("BATCH", {"affine": False}), + }, torch.tensor([1, 1, 1, 1, 1]), torch.tensor([1, 1, 1, 1, 1]), (2, 8, 32, 16, 8), (2, 8, 16, 8, 4), ], [ - {"c_prev": 8, "c": 8, "rate": -1, "arch_code_c": [1, 0, 0, 0, 1]}, + {"c_prev": 8, "c": 8, "rate": -1, "arch_code_c": [1, 0, 0, 0, 1], "act_name": "RELU", "norm_name": "INSTANCE"}, torch.tensor([1, 0, 0, 0, 1]), torch.tensor([0.2, 0.2, 0.2, 0.2, 0.2]), (2, 8, 32, 16, 8), @@ -56,12 +63,35 @@ TEST_CASES_2D = [ [ - {"c_prev": 8, "c": 7, "rate": -1, "arch_code_c": [1, 0, 0, 0, 1], "spatial_dims": 2}, + { + "c_prev": 8, + "c": 7, + "rate": -1, + "arch_code_c": [1, 0, 0, 0, 1], + "spatial_dims": 2, + "act_name": "PRELU", + "norm_name": ("BATCH", {"affine": False}), + }, torch.tensor([1, 0]), torch.tensor([0.2, 0.2]), (2, 8, 16, 8), (2, 7, 8, 4), - ] + ], + [ + { + "c_prev": 8, + "c": 8, + "rate": -1, + "arch_code_c": None, + "spatial_dims": 2, + "act_name": "SELU", + "norm_name": "INSTANCE", + }, + torch.tensor([1, 0]), + torch.tensor([0.2, 0.2]), + (2, 8, 16, 8), + (2, 8, 8, 4), + ], ] diff --git a/tests/test_dints_network.py b/tests/test_dints_network.py index 8be5eb7ccd..08e75fab98 100644 --- a/tests/test_dints_network.py +++ b/tests/test_dints_network.py @@ -33,7 +33,7 @@ "in_channels": 1, "num_classes": 3, "act_name": "RELU", - "norm_name": "INSTANCE", + "norm_name": ("INSTANCE", {"affine": True}), "use_downsample": False, "spatial_dims": 3, }, @@ -101,7 +101,7 @@ "in_channels": 1, "num_classes": 4, "act_name": "RELU", - "norm_name": "INSTANCE", + "norm_name": ("INSTANCE", {"affine": True}), "use_downsample": False, "spatial_dims": 2, }, diff --git a/tests/test_divisible_pad.py b/tests/test_divisible_pad.py index f940636fa8..df610c4939 100644 --- a/tests/test_divisible_pad.py +++ b/tests/test_divisible_pad.py @@ -11,43 +11,32 @@ import unittest -import numpy as np -import torch from parameterized import parameterized from monai.transforms import DivisiblePad -from tests.utils import TEST_NDARRAYS +from monai.utils.enums import NumpyPadMode, PytorchPadMode +from tests.padders import PadTest TESTS = [] -for p in TEST_NDARRAYS: - # pad first dim to be divisible by 7, the second unchanged. - TESTS.append([{"k": (7, -1), "mode": "constant"}, p(np.zeros((3, 8, 7))), p(np.zeros((3, 14, 7)))]) +# pad first dim to be divisible by 7, the second unchanged. +TESTS.append([{"k": (7, -1)}, (3, 8, 7), (3, 14, 7)]) +# pad all dimensions to be divisible by 5 +TESTS.append([{"k": 5, "method": "end"}, (3, 10, 5, 17), (3, 10, 5, 20)]) - # pad all dimensions to be divisible by 5 - TESTS.append( - [{"k": 5, "mode": "constant", "method": "end"}, p(np.zeros((3, 10, 5, 17))), p(np.zeros((3, 10, 5, 20)))] - ) +class TestDivisiblePad(PadTest): + Padder = DivisiblePad -class TestDivisiblePad(unittest.TestCase): @parameterized.expand(TESTS) - def test_pad_shape(self, input_param, input_data, expected_val): - padder = DivisiblePad(**input_param) - result = padder(input_data) - self.assertAlmostEqual(result.shape, expected_val.shape) - result = padder(input_data, mode=input_param["mode"]) - self.assertAlmostEqual(result.shape, expected_val.shape) + def test_pad(self, input_param, input_shape, expected_shape): + modes = ["constant", NumpyPadMode.CONSTANT, PytorchPadMode.CONSTANT] + self.pad_test(input_param, input_shape, expected_shape, modes) def test_pad_kwargs(self): - for p in TEST_NDARRAYS: - input_data = p(np.zeros((3, 8, 4))) - if isinstance(input_data, np.ndarray): - result = DivisiblePad(k=5, mode="constant", constant_values=((0, 0), (1, 1), (2, 2)))(input_data) - np.testing.assert_allclose(result[:, :1, :4], np.ones((3, 1, 4)), rtol=1e-7, atol=0) - else: - result = DivisiblePad(k=5, mode="constant", value=2)(input_data).cpu() - torch.testing.assert_allclose(result[:, :, 4:5], np.ones((3, 10, 1)) + 1, rtol=1e-7, atol=0) + kwargs = {"k": 5, "method": "end"} + unchanged_slices = [slice(None), slice(None, 8), slice(None, 4)] + self.pad_test_kwargs(unchanged_slices, **kwargs) if __name__ == "__main__": diff --git a/tests/test_divisible_padd.py b/tests/test_divisible_padd.py index 61fe917421..93e5a879f0 100644 --- a/tests/test_divisible_padd.py +++ b/tests/test_divisible_padd.py @@ -11,32 +11,25 @@ import unittest -import numpy as np from parameterized import parameterized from monai.transforms import DivisiblePadd +from monai.utils.enums import NumpyPadMode, PytorchPadMode +from tests.padders import PadTest -TEST_CASE_1 = [ - {"keys": ["img"], "k": [4, 3, 2], "mode": "constant"}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 8, 9, 4)), +TESTS = [ + [{"keys": "img", "k": [4, 3, 2]}, (3, 8, 8, 4), (3, 8, 9, 4)], + [{"keys": "img", "k": 7, "method": "end"}, (3, 8, 7), (3, 14, 7)], ] -TEST_CASE_2 = [ - {"keys": ["img"], "k": 7, "mode": "constant", "method": "end"}, - {"img": np.zeros((3, 8, 7))}, - np.zeros((3, 14, 7)), -] - -TEST_CASE_3 = [{"keys": ["img"], "k": 0, "mode": {"constant"}}, {"img": np.zeros((3, 8))}, np.zeros((3, 8))] +class TestDivisiblePadd(PadTest): + Padder = DivisiblePadd -class TestDivisiblePadd(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) - def test_pad_shape(self, input_param, input_data, expected_val): - padder = DivisiblePadd(**input_param) - result = padder(input_data) - np.testing.assert_allclose(result["img"], expected_val) + @parameterized.expand(TESTS) + def test_pad(self, input_param, input_shape, expected_shape): + modes = ["constant", NumpyPadMode.CONSTANT, PytorchPadMode.CONSTANT, "edge", NumpyPadMode.EDGE] + self.pad_test(input_param, input_shape, expected_shape, modes) if __name__ == "__main__": diff --git a/tests/test_drop_path.py b/tests/test_drop_path.py new file mode 100644 index 0000000000..f8ea454228 --- /dev/null +++ b/tests/test_drop_path.py @@ -0,0 +1,43 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks.layers import DropPath + +TEST_CASES = [ + [{"drop_prob": 0.0, "scale_by_keep": True}, (1, 8, 8)], + [{"drop_prob": 0.7, "scale_by_keep": False}, (2, 16, 16, 16)], + [{"drop_prob": 0.3, "scale_by_keep": True}, (6, 16, 12)], +] + +TEST_ERRORS = [[{"drop_prob": 2, "scale_by_keep": False}, (1, 24, 6)]] + + +class TestDropPath(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_shape(self, input_param, input_shape): + im = torch.rand(input_shape) + dr_path = DropPath(**input_param) + out = dr_path(im) + self.assertEqual(out.shape, input_shape) + + @parameterized.expand(TEST_ERRORS) + def test_ill_arg(self, input_param, input_shape): + with self.assertRaises(ValueError): + DropPath(**input_param) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_dynunet.py b/tests/test_dynunet.py index 36ac9d0309..01c58ea788 100644 --- a/tests/test_dynunet.py +++ b/tests/test_dynunet.py @@ -17,7 +17,8 @@ from monai.networks import eval_mode from monai.networks.nets import DynUNet -from tests.utils import test_script_save +from monai.utils.module import pytorch_after +from tests.utils import skip_if_no_cuda, skip_if_windows, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -104,9 +105,11 @@ class TestDynUNet(unittest.TestCase): - @parameterized.expand(TEST_CASE_DYNUNET_2D + TEST_CASE_DYNUNET_3D) + @parameterized.expand(TEST_CASE_DYNUNET_3D) def test_shape(self, input_param, input_shape, expected_shape): net = DynUNet(**input_param).to(device) + if "alphadropout" in input_param.get("dropout"): + self.assertTrue(any(isinstance(x, torch.nn.AlphaDropout) for x in net.modules())) with eval_mode(net): result = net(torch.randn(input_shape).to(device)) self.assertEqual(result.shape, expected_shape) @@ -118,6 +121,36 @@ def test_script(self): test_script_save(net, test_data) +@skip_if_no_cuda +@skip_if_windows +class TestDynUNetWithInstanceNorm3dNVFuser(unittest.TestCase): + @parameterized.expand([TEST_CASE_DYNUNET_3D[0]]) + def test_consistency(self, input_param, input_shape, _): + for eps in [1e-4, 1e-5]: + for momentum in [0.1, 0.01]: + for affine in [True, False]: + norm_param = {"eps": eps, "momentum": momentum, "affine": affine} + input_param["norm_name"] = ("instance", norm_param) + input_param_fuser = input_param.copy() + input_param_fuser["norm_name"] = ("instance_nvfuser", norm_param) + for memory_format in [torch.contiguous_format, torch.channels_last_3d]: + net = DynUNet(**input_param).to("cuda:0", memory_format=memory_format) + net_fuser = DynUNet(**input_param_fuser).to("cuda:0", memory_format=memory_format) + net_fuser.load_state_dict(net.state_dict()) + + input_tensor = torch.randn(input_shape).to("cuda:0", memory_format=memory_format) + with eval_mode(net): + result = net(input_tensor) + with eval_mode(net_fuser): + result_fuser = net_fuser(input_tensor) + + # torch.testing.assert_allclose() is deprecated since 1.12 and will be removed in 1.14 + if pytorch_after(1, 12): + torch.testing.assert_close(result, result_fuser) + else: + torch.testing.assert_allclose(result, result_fuser) + + class TestDynUNetDeepSupervision(unittest.TestCase): @parameterized.expand(TEST_CASE_DEEP_SUPERVISION) def test_shape(self, input_param, input_shape, expected_shape): diff --git a/tests/test_efficientnet.py b/tests/test_efficientnet.py index a2a5e30750..d56f901af7 100644 --- a/tests/test_efficientnet.py +++ b/tests/test_efficientnet.py @@ -367,7 +367,8 @@ def test_func_get_efficientnet_input_shape(self): self.assertEqual(result_shape, expected_shape) def test_script(self): - net = EfficientNetBN(model_name="efficientnet-b0", spatial_dims=2, in_channels=3, num_classes=1000) + with skip_if_downloading_fails(): + net = EfficientNetBN(model_name="efficientnet-b0", spatial_dims=2, in_channels=3, num_classes=1000) net.set_swish(memory_efficient=False) # at the moment custom memory efficient swish is not exportable with jit test_data = torch.randn(1, 3, 224, 224) test_script_save(net, test_data) diff --git a/tests/test_ensure_channel_first.py b/tests/test_ensure_channel_first.py index dd6168ec75..1cb5ac6dec 100644 --- a/tests/test_ensure_channel_first.py +++ b/tests/test_ensure_channel_first.py @@ -16,30 +16,27 @@ import itk import nibabel as nib import numpy as np +import torch from parameterized import parameterized from PIL import Image from monai.data import ITKReader +from monai.data.meta_tensor import MetaTensor from monai.transforms import EnsureChannelFirst, LoadImage -from tests.utils import TEST_NDARRAYS -TEST_CASE_1 = [{"image_only": False}, ["test_image.nii.gz"], None] +TEST_CASE_1 = [{}, ["test_image.nii.gz"], None] -TEST_CASE_2 = [{"image_only": False}, ["test_image.nii.gz"], -1] +TEST_CASE_2 = [{}, ["test_image.nii.gz"], -1] -TEST_CASE_3 = [{"image_only": False}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], None] +TEST_CASE_3 = [{}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], None] -TEST_CASE_4 = [{"reader": ITKReader(), "image_only": False}, ["test_image.nii.gz"], None] +TEST_CASE_4 = [{"reader": ITKReader()}, ["test_image.nii.gz"], None] -TEST_CASE_5 = [{"reader": ITKReader(), "image_only": False}, ["test_image.nii.gz"], -1] +TEST_CASE_5 = [{"reader": ITKReader()}, ["test_image.nii.gz"], -1] -TEST_CASE_6 = [ - {"reader": ITKReader(), "image_only": False}, - ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], - None, -] +TEST_CASE_6 = [{"reader": ITKReader()}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], None] -TEST_CASE_7 = [{"image_only": False, "reader": ITKReader(pixel_type=itk.UC)}, "tests/testing_data/CT_DICOM", None] +TEST_CASE_7 = [{"reader": ITKReader(pixel_type=itk.UC)}, "tests/testing_data/CT_DICOM", None] class TestEnsureChannelFirst(unittest.TestCase): @@ -54,15 +51,15 @@ def test_load_nifti(self, input_param, filenames, original_channel_dim): for i, name in enumerate(filenames): filenames[i] = os.path.join(tempdir, name) nib.save(nib.Nifti1Image(test_image, np.eye(4)), filenames[i]) - for p in TEST_NDARRAYS: - result, header = LoadImage(**input_param)(filenames) - result = EnsureChannelFirst()(p(result), header) - self.assertEqual(result.shape[0], len(filenames)) + + result = LoadImage(image_only=True, **input_param)(filenames) + result = EnsureChannelFirst()(result) + self.assertEqual(result.shape[0], len(filenames)) @parameterized.expand([TEST_CASE_7]) - def test_itk_dicom_series_reader(self, input_param, filenames, original_channel_dim): - result, header = LoadImage(**input_param)(filenames) - result = EnsureChannelFirst()(result, header) + def test_itk_dicom_series_reader(self, input_param, filenames, _): + result = LoadImage(image_only=True, **input_param)(filenames) + result = EnsureChannelFirst()(result) self.assertEqual(result.shape[0], 1) def test_load_png(self): @@ -71,17 +68,20 @@ def test_load_png(self): with tempfile.TemporaryDirectory() as tempdir: filename = os.path.join(tempdir, "test_image.png") Image.fromarray(test_image.astype("uint8")).save(filename) - result, header = LoadImage(image_only=False)(filename) - result = EnsureChannelFirst()(result, header) + result = LoadImage(image_only=True)(filename) + result = EnsureChannelFirst()(result) self.assertEqual(result.shape[0], 3) def test_check(self): + im = torch.zeros(1, 2, 3) + with self.assertRaises(ValueError): # not MetaTensor + EnsureChannelFirst()(im) with self.assertRaises(ValueError): # no meta - EnsureChannelFirst()(np.zeros((1, 2, 3)), None) + EnsureChannelFirst()(MetaTensor(im)) with self.assertRaises(ValueError): # no meta channel - EnsureChannelFirst()(np.zeros((1, 2, 3)), {"original_channel_dim": None}) - EnsureChannelFirst(strict_check=False)(np.zeros((1, 2, 3)), None) - EnsureChannelFirst(strict_check=False)(np.zeros((1, 2, 3)), {"original_channel_dim": None}) + EnsureChannelFirst()(MetaTensor(im, meta={"original_channel_dim": None})) + EnsureChannelFirst(strict_check=False)(im) + EnsureChannelFirst(strict_check=False)(MetaTensor(im, meta={"original_channel_dim": None})) if __name__ == "__main__": diff --git a/tests/test_ensure_channel_firstd.py b/tests/test_ensure_channel_firstd.py index 7f1a57a207..8525939f59 100644 --- a/tests/test_ensure_channel_firstd.py +++ b/tests/test_ensure_channel_firstd.py @@ -15,12 +15,12 @@ import nibabel as nib import numpy as np +import torch from parameterized import parameterized from PIL import Image +from monai.data.meta_tensor import MetaTensor from monai.transforms import EnsureChannelFirstd, LoadImaged -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS TEST_CASE_1 = [{"keys": "img"}, ["test_image.nii.gz"], None] @@ -41,11 +41,9 @@ def test_load_nifti(self, input_param, filenames, original_channel_dim): for i, name in enumerate(filenames): filenames[i] = os.path.join(tempdir, name) nib.save(nib.Nifti1Image(test_image, np.eye(4)), filenames[i]) - for p in TEST_NDARRAYS: - result = LoadImaged(**input_param)({"img": filenames}) - result["img"] = p(result["img"]) - result = EnsureChannelFirstd(**input_param)(result) - self.assertEqual(result["img"].shape[0], len(filenames)) + result = LoadImaged(**input_param)({"img": filenames}) + result = EnsureChannelFirstd(**input_param)(result) + self.assertEqual(result["img"].shape[0], len(filenames)) def test_load_png(self): spatial_size = (256, 256, 3) @@ -58,16 +56,13 @@ def test_load_png(self): self.assertEqual(result["img"].shape[0], 3) def test_exceptions(self): + im = torch.zeros((1, 2, 3)) with self.assertRaises(ValueError): # no meta - EnsureChannelFirstd("img")({"img": np.zeros((1, 2, 3)), PostFix.meta("img"): None}) + EnsureChannelFirstd("img")({"img": im}) with self.assertRaises(ValueError): # no meta channel - EnsureChannelFirstd("img")( - {"img": np.zeros((1, 2, 3)), PostFix.meta("img"): {"original_channel_dim": None}} - ) - EnsureChannelFirstd("img", strict_check=False)({"img": np.zeros((1, 2, 3)), PostFix.meta("img"): None}) - EnsureChannelFirstd("img", strict_check=False)( - {"img": np.zeros((1, 2, 3)), PostFix.meta("img"): {"original_channel_dim": None}} - ) + EnsureChannelFirstd("img")({"img": MetaTensor(im, meta={"original_channel_dim": None})}) + EnsureChannelFirstd("img", strict_check=False)({"img": im}) + EnsureChannelFirstd("img", strict_check=False)({"img": MetaTensor(im, meta={"original_channel_dim": None})}) if __name__ == "__main__": diff --git a/tests/test_ensure_tuple.py b/tests/test_ensure_tuple.py new file mode 100644 index 0000000000..ea580871da --- /dev/null +++ b/tests/test_ensure_tuple.py @@ -0,0 +1,52 @@ +# 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. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.utils.misc import ensure_tuple +from tests.utils import assert_allclose + +TESTS = [ + ["test", ("test",)], + [["test1", "test2"], ("test1", "test2")], + [123, (123,)], + [(1, [2], 3), (1, [2], 3)], + [(1, 2, 3), (1, 2, 3), True], + [np.array([1, 2]), (np.array([1, 2]),), True], + [np.array([1, 2]), (1, 2), False], + [torch.tensor([1, 2]), (torch.tensor([1, 2]),), True], + [np.array([]), (np.array([]),)], + [torch.tensor([]), (torch.tensor([]),)], + [np.array(123), (np.array(123),), True], + [torch.tensor(123), (torch.tensor(123),)], +] + + +class TestEnsureTuple(unittest.TestCase): + @parameterized.expand(TESTS) + def test_value(self, input, expected_value, wrap_array=False): + result = ensure_tuple(input, wrap_array) + + self.assertTrue(isinstance(result, tuple)) + if isinstance(input, (np.ndarray, torch.Tensor)): + for i, j in zip(result, expected_value): + assert_allclose(i, j) + else: + self.assertTupleEqual(result, expected_value) + + +if __name__ == "__main__": + + unittest.main() diff --git a/tests/test_ensure_type.py b/tests/test_ensure_type.py index f8a6ee30ff..55423838b8 100644 --- a/tests/test_ensure_type.py +++ b/tests/test_ensure_type.py @@ -14,6 +14,7 @@ import numpy as np import torch +from monai.data import MetaTensor from monai.transforms import EnsureType from tests.utils import assert_allclose @@ -59,9 +60,9 @@ def test_string(self): def test_list_tuple(self): for dtype in ("tensor", "numpy"): - result = EnsureType(data_type=dtype, wrap_sequence=False)([[1, 2], [3, 4]]) + result = EnsureType(data_type=dtype, wrap_sequence=False, track_meta=True)([[1, 2], [3, 4]]) self.assertTrue(isinstance(result, list)) - self.assertTrue(isinstance(result[0][1], torch.Tensor if dtype == "tensor" else np.ndarray)) + self.assertTrue(isinstance(result[0][1], MetaTensor if dtype == "tensor" else np.ndarray)) torch.testing.assert_allclose(result[1][0], torch.as_tensor(3)) # tuple of numpy arrays result = EnsureType(data_type=dtype, wrap_sequence=False)((np.array([1, 2]), np.array([3, 4]))) @@ -77,7 +78,7 @@ def test_dict(self): "extra": None, } for dtype in ("tensor", "numpy"): - result = EnsureType(data_type=dtype)(test_data) + result = EnsureType(data_type=dtype, track_meta=False)(test_data) self.assertTrue(isinstance(result, dict)) self.assertTrue(isinstance(result["img"], torch.Tensor if dtype == "tensor" else np.ndarray)) torch.testing.assert_allclose(result["img"], torch.as_tensor([1.0, 2.0])) diff --git a/tests/test_ensure_typed.py b/tests/test_ensure_typed.py index cadab9bd56..d57170e2a6 100644 --- a/tests/test_ensure_typed.py +++ b/tests/test_ensure_typed.py @@ -14,6 +14,7 @@ import numpy as np import torch +from monai.data import MetaTensor from monai.transforms import EnsureTyped from tests.utils import assert_allclose @@ -61,9 +62,11 @@ def test_string(self): def test_list_tuple(self): for dtype in ("tensor", "numpy"): - result = EnsureTyped(keys="data", data_type=dtype, wrap_sequence=False)({"data": [[1, 2], [3, 4]]})["data"] + result = EnsureTyped(keys="data", data_type=dtype, wrap_sequence=False, track_meta=True)( + {"data": [[1, 2], [3, 4]]} + )["data"] self.assertTrue(isinstance(result, list)) - self.assertTrue(isinstance(result[0][1], torch.Tensor if dtype == "tensor" else np.ndarray)) + self.assertTrue(isinstance(result[0][1], MetaTensor if dtype == "tensor" else np.ndarray)) torch.testing.assert_allclose(result[1][0], torch.as_tensor(3)) # tuple of numpy arrays result = EnsureTyped(keys="data", data_type=dtype, wrap_sequence=False)( diff --git a/tests/test_fft_utils.py b/tests/test_fft_utils.py new file mode 100644 index 0000000000..d5e3a22eaa --- /dev/null +++ b/tests/test_fft_utils.py @@ -0,0 +1,63 @@ +# 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. + +import unittest + +from parameterized import parameterized + +from monai.data.fft_utils import fftn_centered, ifftn_centered +from tests.utils import TEST_NDARRAYS, assert_allclose + +# +im = [[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]] +res = [ + [[[0.0, 0.0], [0.0, 3.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, 0.0]]] +] +TESTS = [] +for p in TEST_NDARRAYS: + TESTS.append((p(im), p(res))) + +# +TESTS_CONSISTENCY = [] +for p in TEST_NDARRAYS: + TESTS_CONSISTENCY.append(p(im)) + +# +im_complex = [ + [[[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]]] +] +TESTS_CONSISTENCY_COMPLEX = [] +for p in TEST_NDARRAYS: + TESTS_CONSISTENCY_COMPLEX.append(p(im_complex)) + + +class TestFFT(unittest.TestCase): + @parameterized.expand(TESTS) + def test(self, test_data, res_data): + result = fftn_centered(test_data, spatial_dims=2, is_complex=False) + assert_allclose(result, res_data, type_test=True) + + @parameterized.expand(TESTS_CONSISTENCY) + def test_consistency(self, test_data): + result = fftn_centered(test_data, spatial_dims=2, is_complex=False) + result = ifftn_centered(result, spatial_dims=2, is_complex=True) + result = (result[..., 0] ** 2 + result[..., 1] ** 2) ** 0.5 + assert_allclose(result, test_data, type_test=False) + + @parameterized.expand(TESTS_CONSISTENCY_COMPLEX) + def test_consistency_complex(self, test_data): + result = fftn_centered(test_data, spatial_dims=2) + result = ifftn_centered(result, spatial_dims=2) + assert_allclose(result, test_data, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_fill_holes.py b/tests/test_fill_holes.py index 9f9dc1fc2e..4292ff3a22 100644 --- a/tests/test_fill_holes.py +++ b/tests/test_fill_holes.py @@ -192,10 +192,6 @@ TEST_CASE_22, ] -ITEST_CASE_1 = ["invalid_image_data_type", {}, [[[[1, 1, 1]]]], NotImplementedError] - -INVALID_CASES = [ITEST_CASE_1] - class TestFillHoles(unittest.TestCase): @parameterized.expand(VALID_CASES) @@ -203,16 +199,7 @@ def test_correct_results(self, _, args, input_image, expected): converter = FillHoles(**args) for p in TEST_NDARRAYS: result = converter(p(clone(input_image))) - assert_allclose(result, p(expected)) - - @parameterized.expand(INVALID_CASES) - def test_raise_exception(self, _, args, input_image, expected_error): - with self.assertRaises(expected_error): - converter = FillHoles(**args) - if isinstance(input_image, torch.Tensor) and torch.cuda.is_available(): - _ = converter(clone(input_image).cuda()) - else: - _ = converter(clone(input_image)) + assert_allclose(result, p(expected), type_test=False) if __name__ == "__main__": diff --git a/tests/test_fill_holesd.py b/tests/test_fill_holesd.py index f7aa9f6108..fce90fd86a 100644 --- a/tests/test_fill_holesd.py +++ b/tests/test_fill_holesd.py @@ -193,10 +193,6 @@ TEST_CASE_22, ] -ITEST_CASE_1 = ["invalid_image_data_type", {}, [[[[1, 1, 1]]]], NotImplementedError] - -INVALID_CASES = [ITEST_CASE_1] - class TestFillHoles(unittest.TestCase): @parameterized.expand(VALID_CASES) @@ -205,17 +201,7 @@ def test_correct_results(self, _, args, input_image, expected): converter = FillHolesd(keys=key, **args) for p in TEST_NDARRAYS: result = converter({key: p(clone(input_image))})[key] - assert_allclose(result, p(expected)) - - @parameterized.expand(INVALID_CASES) - def test_raise_exception(self, _, args, input_image, expected_error): - key = CommonKeys.IMAGE - with self.assertRaises(expected_error): - converter = FillHolesd(keys=key, **args) - if isinstance(input_image, torch.Tensor) and torch.cuda.is_available(): - _ = converter({key: clone(input_image).cuda()})[key] - else: - _ = converter({key: clone(input_image)})[key] + assert_allclose(result, p(expected), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_flip.py b/tests/test_flip.py index 17cf0d2c39..c5a281b127 100644 --- a/tests/test_flip.py +++ b/tests/test_flip.py @@ -12,15 +12,23 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.transforms import Flip -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_DEVICES, TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion INVALID_CASES = [("wrong_axis", ["s", 1], TypeError), ("not_numbers", "s", TypeError)] VALID_CASES = [("no_axis", None), ("one_axis", 1), ("many_axis", [0, 1]), ("negative_axis", [0, -1])] +TORCH_CASES = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TORCH_CASES.append([[0, 1], torch.zeros((1, 3, 2)), track_meta, *device]) + class TestFlip(NumpyImageTestCase2D): @parameterized.expand(INVALID_CASES) @@ -31,13 +39,29 @@ def test_invalid_inputs(self, _, spatial_axis, raises): @parameterized.expand(VALID_CASES) def test_correct_results(self, _, spatial_axis): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: im = p(self.imt[0]) flip = Flip(spatial_axis=spatial_axis) expected = [np.flip(channel, spatial_axis) for channel in self.imt[0]] expected = np.stack(expected) result = flip(im) - assert_allclose(result, p(expected)) + assert_allclose(result, p(expected), type_test="tensor") + test_local_inversion(flip, result, im) + + @parameterized.expand(TORCH_CASES) + def test_torch(self, init_param, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + img = img.to(device) + xform = Flip(init_param) + res = xform(img) + self.assertEqual(img.shape, res.shape) + if track_meta: + self.assertIsInstance(res, MetaTensor) + else: + self.assertNotIsInstance(res, MetaTensor) + self.assertIsInstance(res, torch.Tensor) + with self.assertRaisesRegex(ValueError, "MetaTensor"): + xform.inverse(res) if __name__ == "__main__": diff --git a/tests/test_flipd.py b/tests/test_flipd.py index 900779f4e0..c97674b83b 100644 --- a/tests/test_flipd.py +++ b/tests/test_flipd.py @@ -12,15 +12,23 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.transforms import Flipd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_DEVICES, TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion INVALID_CASES = [("wrong_axis", ["s", 1], TypeError), ("not_numbers", "s", TypeError)] VALID_CASES = [("no_axis", None), ("one_axis", 1), ("many_axis", [0, 1])] +TORCH_CASES = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TORCH_CASES.append([[0, 1], torch.zeros((1, 3, 2)), track_meta, *device]) + class TestFlipd(NumpyImageTestCase2D): @parameterized.expand(INVALID_CASES) @@ -31,12 +39,29 @@ def test_invalid_cases(self, _, spatial_axis, raises): @parameterized.expand(VALID_CASES) def test_correct_results(self, _, spatial_axis): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: flip = Flipd(keys="img", spatial_axis=spatial_axis) expected = [np.flip(channel, spatial_axis) for channel in self.imt[0]] expected = np.stack(expected) - result = flip({"img": p(self.imt[0])})["img"] - assert_allclose(result, p(expected)) + im = p(self.imt[0]) + result = flip({"img": im})["img"] + assert_allclose(result, p(expected), type_test="tensor") + test_local_inversion(flip, {"img": result}, {"img": im}, "img") + + @parameterized.expand(TORCH_CASES) + def test_torch(self, init_param, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + img = img.to(device) + xform = Flipd("image", init_param) + res = xform({"image": img}) + self.assertEqual(img.shape, res["image"].shape) + if track_meta: + self.assertIsInstance(res["image"], MetaTensor) + else: + self.assertNotIsInstance(res["image"], MetaTensor) + self.assertIsInstance(res["image"], torch.Tensor) + with self.assertRaisesRegex(ValueError, "MetaTensor"): + xform.inverse(res) if __name__ == "__main__": diff --git a/tests/test_focal_loss.py b/tests/test_focal_loss.py index d8a9c8ab5b..6ac23fef36 100644 --- a/tests/test_focal_loss.py +++ b/tests/test_focal_loss.py @@ -17,7 +17,7 @@ from monai.losses import FocalLoss from monai.networks import one_hot -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save class TestFocalLoss(unittest.TestCase): @@ -67,11 +67,8 @@ def test_consistency_with_cross_entropy_2d_no_reduction(self): b = output1.cpu().detach().numpy() error = np.abs(a - b) max_error = np.maximum(error, max_error) - # if np.all(np.abs(a - b) > max_error): - # max_error = np.abs(a - b) assert np.allclose(max_error, 0) - # self.assertAlmostEqual(max_error, 0.0, places=3) def test_consistency_with_cross_entropy_2d_onehot_label(self): """For gamma=0 the focal loss reduces to the cross entropy loss""" @@ -261,7 +258,6 @@ def test_ill_class_weight(self): with self.assertRaisesRegex(ValueError, ""): FocalLoss(include_background=False, weight=(1.0, 1.0, -1.0))(chn_input, chn_target) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = FocalLoss() test_input = torch.ones(2, 2, 8, 8) diff --git a/tests/test_foreground_mask.py b/tests/test_foreground_mask.py new file mode 100644 index 0000000000..160db5bae3 --- /dev/null +++ b/tests/test_foreground_mask.py @@ -0,0 +1,96 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.intensity.array import ForegroundMask +from monai.utils import min_version, optional_import, set_determinism +from tests.utils import TEST_NDARRAYS, assert_allclose + +skimage, has_skimage = optional_import("skimage", "0.19.0", min_version) +set_determinism(1234) + +A = np.random.randint(64, 128, (3, 3, 2)).astype(np.uint8) +A3D = np.random.randint(64, 128, (3, 3, 2, 2)).astype(np.uint8) +B = np.ones_like(A[:1]) +B3D = np.ones_like(A3D[:1]) +MASK = np.pad(B, ((0, 0), (2, 2), (2, 2)), constant_values=0) +MASK3D = np.pad(B3D, ((0, 0), (2, 2), (2, 2), (2, 2)), constant_values=0) +IMAGE1 = np.pad(A, ((0, 0), (2, 2), (2, 2)), constant_values=255) +IMAGE3D = np.pad(A3D, ((0, 0), (2, 2), (2, 2), (2, 2)), constant_values=255) +IMAGE2 = np.copy(IMAGE1) +IMAGE2[0] = 0 +IMAGE3 = np.pad(A, ((0, 0), (2, 2), (2, 2)), constant_values=0) +TEST_CASE_0 = [{}, IMAGE1, MASK] +TEST_CASE_1 = [{"threshold": "otsu"}, IMAGE1, MASK] +TEST_CASE_2 = [{"threshold": "otsu"}, IMAGE2, MASK] +TEST_CASE_3 = [{"threshold": 140}, IMAGE1, MASK] +TEST_CASE_4 = [{"threshold": "otsu", "invert": True}, IMAGE3, MASK] +TEST_CASE_5 = [{"threshold": 0.5}, MASK, np.logical_not(MASK)] +TEST_CASE_6 = [{"threshold": 140}, IMAGE2, np.ones_like(MASK)] +TEST_CASE_7 = [{"threshold": {"R": "otsu", "G": "otsu", "B": "otsu"}}, IMAGE2, MASK] +TEST_CASE_8 = [{"threshold": {"R": 140, "G": "otsu", "B": "otsu"}}, IMAGE2, np.ones_like(MASK)] +TEST_CASE_9 = [{"threshold": {"R": 140, "G": skimage.filters.threshold_otsu, "B": "otsu"}}, IMAGE2, np.ones_like(MASK)] +TEST_CASE_10 = [{"threshold": skimage.filters.threshold_mean}, IMAGE1, MASK] +TEST_CASE_11 = [{"threshold": None, "hsv_threshold": "otsu"}, IMAGE1, np.ones_like(MASK)] +TEST_CASE_12 = [{"threshold": None, "hsv_threshold": {"S": "otsu"}}, IMAGE1, MASK] +TEST_CASE_13 = [{"threshold": 100, "invert": True}, IMAGE1, np.logical_not(MASK)] +TEST_CASE_14 = [{}, IMAGE3D, MASK3D] +TEST_CASE_15 = [{"hsv_threshold": {"S": 0.1}}, IMAGE3D, MASK3D] + +TEST_CASE_ERROR_1 = [{"threshold": None}, IMAGE1] +TEST_CASE_ERROR_2 = [{"threshold": {"K": 1}}, IMAGE1] + +TESTS = [] +for p in TEST_NDARRAYS: + TESTS.append([p, *TEST_CASE_0]) + TESTS.append([p, *TEST_CASE_1]) + TESTS.append([p, *TEST_CASE_2]) + TESTS.append([p, *TEST_CASE_3]) + TESTS.append([p, *TEST_CASE_4]) + TESTS.append([p, *TEST_CASE_5]) + TESTS.append([p, *TEST_CASE_6]) + TESTS.append([p, *TEST_CASE_7]) + TESTS.append([p, *TEST_CASE_8]) + TESTS.append([p, *TEST_CASE_9]) + TESTS.append([p, *TEST_CASE_10]) + TESTS.append([p, *TEST_CASE_11]) + TESTS.append([p, *TEST_CASE_12]) + TESTS.append([p, *TEST_CASE_13]) + TESTS.append([p, *TEST_CASE_14]) + TESTS.append([p, *TEST_CASE_15]) + +TESTS_ERROR = [] +for p in TEST_NDARRAYS: + TESTS_ERROR.append([p, *TEST_CASE_ERROR_1]) + TESTS_ERROR.append([p, *TEST_CASE_ERROR_2]) + + +@unittest.skipUnless(has_skimage, "Requires sci-kit image") +class TestForegroundMask(unittest.TestCase): + @parameterized.expand(TESTS) + def test_foreground_mask(self, in_type, arguments, image, mask): + input_image = in_type(image) + result = ForegroundMask(**arguments)(input_image) + assert_allclose(result, mask, type_test="tensor") + + @parameterized.expand(TESTS_ERROR) + def test_foreground_mask_error(self, in_type, arguments, image): + input_image = in_type(image) + with self.assertRaises(ValueError): + ForegroundMask(**arguments)(input_image) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_foreground_maskd.py b/tests/test_foreground_maskd.py new file mode 100644 index 0000000000..3c8aa08d7f --- /dev/null +++ b/tests/test_foreground_maskd.py @@ -0,0 +1,104 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.intensity.dictionary import ForegroundMaskd +from monai.utils import min_version, optional_import, set_determinism +from tests.utils import TEST_NDARRAYS, assert_allclose + +skimage, has_skimage = optional_import("skimage", "0.19.0", min_version) +set_determinism(1234) + +A = np.random.randint(64, 128, (3, 3, 2)).astype(np.uint8) +A3D = np.random.randint(64, 128, (3, 3, 2, 2)).astype(np.uint8) +B = np.ones_like(A[:1]) +B3D = np.ones_like(A3D[:1]) +MASK = np.pad(B, ((0, 0), (2, 2), (2, 2)), constant_values=0) +MASK3D = np.pad(B3D, ((0, 0), (2, 2), (2, 2), (2, 2)), constant_values=0) +IMAGE1 = np.pad(A, ((0, 0), (2, 2), (2, 2)), constant_values=255) +IMAGE3D = np.pad(A3D, ((0, 0), (2, 2), (2, 2), (2, 2)), constant_values=255) +IMAGE2 = np.copy(IMAGE1) +IMAGE2[0] = 0 +IMAGE3 = np.pad(A, ((0, 0), (2, 2), (2, 2)), constant_values=0) +TEST_CASE_0 = [{"keys": "image"}, {"image": IMAGE1}, MASK] +TEST_CASE_1 = [{"keys": "image", "threshold": "otsu"}, {"image": IMAGE1}, MASK] +TEST_CASE_2 = [{"keys": "image", "threshold": "otsu"}, {"image": IMAGE2}, MASK] +TEST_CASE_3 = [{"keys": "image", "threshold": 140}, {"image": IMAGE1}, MASK] +TEST_CASE_4 = [{"keys": "image", "threshold": "otsu", "invert": True}, {"image": IMAGE3}, MASK] +TEST_CASE_5 = [{"keys": "image", "threshold": 0.5}, {"image": MASK}, np.logical_not(MASK)] +TEST_CASE_6 = [{"keys": "image", "threshold": 140}, {"image": IMAGE2}, np.ones_like(MASK)] +TEST_CASE_7 = [{"keys": "image", "threshold": {"R": "otsu", "G": "otsu", "B": "otsu"}}, {"image": IMAGE2}, MASK] +TEST_CASE_8 = [ + {"keys": "image", "threshold": {"R": 140, "G": "otsu", "B": "otsu"}}, + {"image": IMAGE2}, + np.ones_like(MASK), +] +TEST_CASE_9 = [ + {"keys": "image", "threshold": {"R": 140, "G": skimage.filters.threshold_otsu, "B": "otsu"}}, + {"image": IMAGE2}, + np.ones_like(MASK), +] +TEST_CASE_10 = [{"keys": "image", "threshold": skimage.filters.threshold_mean}, {"image": IMAGE1}, MASK] +TEST_CASE_11 = [{"keys": "image", "threshold": None, "hsv_threshold": "otsu"}, {"image": IMAGE1}, np.ones_like(MASK)] +TEST_CASE_12 = [{"keys": "image", "threshold": None, "hsv_threshold": {"S": "otsu"}}, {"image": IMAGE1}, MASK] +TEST_CASE_13 = [{"keys": "image", "threshold": 100, "invert": True}, {"image": IMAGE1}, np.logical_not(MASK)] +TEST_CASE_14 = [{"keys": "image"}, {"image": IMAGE3D}, MASK3D] +TEST_CASE_15 = [{"keys": "image", "hsv_threshold": {"S": 0.1}}, {"image": IMAGE3D}, MASK3D] + +TEST_CASE_ERROR_1 = [{"keys": "image", "threshold": None}, {"image": IMAGE1}] +TEST_CASE_ERROR_2 = [{"keys": "image", "threshold": {"K": 1}}, {"image": IMAGE1}] + +TESTS = [] +for p in TEST_NDARRAYS: + TESTS.append([p, *TEST_CASE_0]) + TESTS.append([p, *TEST_CASE_1]) + TESTS.append([p, *TEST_CASE_2]) + TESTS.append([p, *TEST_CASE_3]) + TESTS.append([p, *TEST_CASE_4]) + TESTS.append([p, *TEST_CASE_5]) + TESTS.append([p, *TEST_CASE_6]) + TESTS.append([p, *TEST_CASE_7]) + TESTS.append([p, *TEST_CASE_8]) + TESTS.append([p, *TEST_CASE_9]) + TESTS.append([p, *TEST_CASE_10]) + TESTS.append([p, *TEST_CASE_11]) + TESTS.append([p, *TEST_CASE_12]) + TESTS.append([p, *TEST_CASE_13]) + TESTS.append([p, *TEST_CASE_14]) + TESTS.append([p, *TEST_CASE_15]) + +TESTS_ERROR = [] +for p in TEST_NDARRAYS: + TESTS_ERROR.append([p, *TEST_CASE_ERROR_1]) + TESTS_ERROR.append([p, *TEST_CASE_ERROR_2]) + + +@unittest.skipUnless(has_skimage, "Requires sci-kit image") +class TestForegroundMaskd(unittest.TestCase): + @parameterized.expand(TESTS) + def test_foreground_mask(self, in_type, arguments, data_dict, mask): + data_dict[arguments["keys"]] = in_type(data_dict[arguments["keys"]]) + result = ForegroundMaskd(**arguments)(data_dict)[arguments["keys"]] + assert_allclose(result, mask, type_test=False) + + @parameterized.expand(TESTS_ERROR) + def test_foreground_mask_error(self, in_type, arguments, data_dict): + data_dict[arguments["keys"]] = in_type(data_dict[arguments["keys"]]) + with self.assertRaises(ValueError): + ForegroundMaskd(**arguments)(data_dict)[arguments["keys"]] + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_fourier.py b/tests/test_fourier.py index e1b5f3089d..b500f266d7 100644 --- a/tests/test_fourier.py +++ b/tests/test_fourier.py @@ -21,8 +21,6 @@ from tests.utils import SkipIfBeforePyTorchVersion, SkipIfNoModule TEST_CASES = [((128, 64),), ((64, 48, 80),)] -# for shape in ((128, 64), (64, 48, 80)): -# TEST_CASES.append(shape) @SkipIfBeforePyTorchVersion((1, 8)) diff --git a/tests/test_fpn_block.py b/tests/test_fpn_block.py new file mode 100644 index 0000000000..420fd04367 --- /dev/null +++ b/tests/test_fpn_block.py @@ -0,0 +1,84 @@ +# 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. + +import unittest +from collections import OrderedDict + +import torch +from parameterized import parameterized + +from monai.networks.blocks.backbone_fpn_utils import _resnet_fpn_extractor +from monai.networks.blocks.feature_pyramid_network import FeaturePyramidNetwork +from monai.networks.nets.resnet import resnet50 +from monai.utils import optional_import +from tests.utils import test_script_save + +_, has_torchvision = optional_import("torchvision") + +TEST_CASES = [ + [ + {"spatial_dims": 3, "in_channels_list": [32, 64], "out_channels": 6}, + ((7, 32, 16, 32, 64), (7, 64, 8, 16, 32)), + ((7, 6, 16, 32, 64), (7, 6, 8, 16, 32)), + ], + [ + {"spatial_dims": 2, "in_channels_list": [32, 64], "out_channels": 6}, + ((7, 32, 16, 32), (7, 64, 8, 16)), + ((7, 6, 16, 32), (7, 6, 8, 16)), + ], +] + +TEST_CASES2 = [ + [{"spatial_dims": 3, "returned_layers": [1]}, (7, 3, 32, 64, 32), ((7, 256, 16, 32, 16), (7, 256, 8, 16, 8))] +] + + +class TestFPNBlock(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_fpn_block(self, input_param, input_shape, expected_shape): + net = FeaturePyramidNetwork(**input_param) + data = OrderedDict() + data["feat0"] = torch.rand(input_shape[0]) + data["feat1"] = torch.rand(input_shape[1]) + result = net(data) + self.assertEqual(result["feat0"].shape, expected_shape[0]) + self.assertEqual(result["feat1"].shape, expected_shape[1]) + + @parameterized.expand(TEST_CASES) + def test_script(self, input_param, input_shape, expected_shape): + # test whether support torchscript + net = FeaturePyramidNetwork(**input_param) + data = OrderedDict() + data["feat0"] = torch.rand(input_shape[0]) + data["feat1"] = torch.rand(input_shape[1]) + test_script_save(net, data) + + +@unittest.skipUnless(has_torchvision, "Requires torchvision") +class TestFPN(unittest.TestCase): + @parameterized.expand(TEST_CASES2) + def test_fpn(self, input_param, input_shape, expected_shape): + net = _resnet_fpn_extractor(backbone=resnet50(), spatial_dims=input_param["spatial_dims"], returned_layers=[1]) + data = torch.rand(input_shape) + result = net(data) + self.assertEqual(result["0"].shape, expected_shape[0]) + self.assertEqual(result["pool"].shape, expected_shape[1]) + + @parameterized.expand(TEST_CASES2) + def test_script(self, input_param, input_shape, expected_shape): + # test whether support torchscript + net = _resnet_fpn_extractor(backbone=resnet50(), spatial_dims=input_param["spatial_dims"], returned_layers=[1]) + data = torch.rand(input_shape) + test_script_save(net, data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_gaussian_filter.py b/tests/test_gaussian_filter.py index 9d76e44cec..c4ffe56896 100644 --- a/tests/test_gaussian_filter.py +++ b/tests/test_gaussian_filter.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.networks.layers import GaussianFilter -from tests.utils import SkipIfBeforePyTorchVersion, skip_if_quick +from tests.utils import skip_if_quick TEST_CASES = [[{"type": "erf", "gt": 2.0}], [{"type": "scalespace", "gt": 3.0}], [{"type": "sampled", "gt": 5.0}]] TEST_CASES_GPU = [[{"type": "erf", "gt": 0.8, "device": "cuda"}], [{"type": "sampled", "gt": 5.0, "device": "cuda"}]] @@ -82,7 +82,6 @@ def code_to_run(self, input_args): ) @parameterized.expand(TEST_CASES + TEST_CASES_GPU + TEST_CASES_3d) - @SkipIfBeforePyTorchVersion((1, 7)) def test_train_quick(self, input_args): self.code_to_run(input_args) diff --git a/tests/test_gaussian_sharpen.py b/tests/test_gaussian_sharpen.py index 547febdfaf..af36e7c03d 100644 --- a/tests/test_gaussian_sharpen.py +++ b/tests/test_gaussian_sharpen.py @@ -83,7 +83,7 @@ class TestGaussianSharpen(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = GaussianSharpen(**argments)(image) - assert_allclose(result, expected_data, atol=0, rtol=1e-4, type_test=False) + assert_allclose(result, expected_data, atol=0, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_gaussian_sharpend.py b/tests/test_gaussian_sharpend.py index d9ef503532..14339fff26 100644 --- a/tests/test_gaussian_sharpend.py +++ b/tests/test_gaussian_sharpend.py @@ -83,7 +83,7 @@ class TestGaussianSharpend(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = GaussianSharpend(**argments)(image) - assert_allclose(result["img"], expected_data, rtol=1e-4, type_test=False) + assert_allclose(result["img"], expected_data, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_gaussian_smooth.py b/tests/test_gaussian_smooth.py index 53f2fc396b..032a60caad 100644 --- a/tests/test_gaussian_smooth.py +++ b/tests/test_gaussian_smooth.py @@ -87,7 +87,7 @@ class TestGaussianSmooth(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = GaussianSmooth(**argments)(image) - assert_allclose(result, expected_data, atol=0, rtol=1e-4, type_test=False) + assert_allclose(result, expected_data, atol=1e-4, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_gaussian_smoothd.py b/tests/test_gaussian_smoothd.py index 839bac81fe..8f5465c848 100644 --- a/tests/test_gaussian_smoothd.py +++ b/tests/test_gaussian_smoothd.py @@ -87,7 +87,7 @@ class TestGaussianSmoothd(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = GaussianSmoothd(**argments)(image) - assert_allclose(result["img"], expected_data, rtol=1e-4, type_test=False) + assert_allclose(result["img"], expected_data, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_generalized_dice_focal_loss.py b/tests/test_generalized_dice_focal_loss.py new file mode 100644 index 0000000000..ef8661c88d --- /dev/null +++ b/tests/test_generalized_dice_focal_loss.py @@ -0,0 +1,73 @@ +# 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. + +import unittest + +import numpy as np +import torch + +from monai.losses import FocalLoss, GeneralizedDiceFocalLoss, GeneralizedDiceLoss +from tests.utils import test_script_save + + +class TestGeneralizedDiceFocalLoss(unittest.TestCase): + def test_result_onehot_target_include_bg(self): + size = [3, 3, 5, 5] + label = torch.randint(low=0, high=2, size=size) + pred = torch.randn(size) + for reduction in ["sum", "mean", "none"]: + common_params = {"include_background": True, "to_onehot_y": False, "reduction": reduction} + for focal_weight in [None, torch.tensor([1.0, 1.0, 2.0]), (3, 2.0, 1)]: + for lambda_focal in [0.5, 1.0, 1.5]: + generalized_dice_focal = GeneralizedDiceFocalLoss( + focal_weight=focal_weight, gamma=1.0, lambda_focal=lambda_focal, **common_params + ) + generalized_dice = GeneralizedDiceLoss(**common_params) + focal = FocalLoss(weight=focal_weight, gamma=1.0, **common_params) + result = generalized_dice_focal(pred, label) + expected_val = generalized_dice(pred, label) + lambda_focal * focal(pred, label) + np.testing.assert_allclose(result, expected_val) + + def test_result_no_onehot_no_bg(self): + size = [3, 3, 5, 5] + label = torch.randint(low=0, high=2, size=size) + label = torch.argmax(label, dim=1, keepdim=True) + pred = torch.randn(size) + for reduction in ["sum", "mean", "none"]: + common_params = {"include_background": False, "to_onehot_y": True, "reduction": reduction} + for focal_weight in [2.0, torch.tensor([1.0, 2.0]), (2.0, 1)]: + for lambda_focal in [0.5, 1.0, 1.5]: + generalized_dice_focal = GeneralizedDiceFocalLoss( + focal_weight=focal_weight, lambda_focal=lambda_focal, **common_params + ) + generalized_dice = GeneralizedDiceLoss(**common_params) + focal = FocalLoss(weight=focal_weight, **common_params) + result = generalized_dice_focal(pred, label) + expected_val = generalized_dice(pred, label) + lambda_focal * focal(pred, label) + np.testing.assert_allclose(result, expected_val) + + def test_ill_shape(self): + loss = GeneralizedDiceFocalLoss() + with self.assertRaisesRegex(ValueError, ""): + loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3))) + + def test_ill_lambda(self): + with self.assertRaisesRegex(ValueError, ""): + GeneralizedDiceFocalLoss(lambda_gdl=-1.0) + + def test_script(self): + loss = GeneralizedDiceFocalLoss() + test_input = torch.ones(2, 1, 8, 8) + test_script_save(loss, test_input, test_input) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_generalized_dice_loss.py b/tests/test_generalized_dice_loss.py index fa301201e4..81f8f4c0b0 100644 --- a/tests/test_generalized_dice_loss.py +++ b/tests/test_generalized_dice_loss.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.losses import GeneralizedDiceLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -193,7 +193,6 @@ def test_batch(self): loss = generalized_dice_loss(prediction, target) self.assertNotEqual(loss.grad_fn, None) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = GeneralizedDiceLoss() test_input = torch.ones(2, 1, 8, 8) diff --git a/tests/test_generalized_wasserstein_dice_loss.py b/tests/test_generalized_wasserstein_dice_loss.py index 2c33d365f4..49a5aa0556 100644 --- a/tests/test_generalized_wasserstein_dice_loss.py +++ b/tests/test_generalized_wasserstein_dice_loss.py @@ -18,7 +18,7 @@ import torch.optim as optim from monai.losses import GeneralizedWassersteinDiceLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save class TestGeneralizedWassersteinDiceLoss(unittest.TestCase): @@ -216,7 +216,6 @@ def forward(self, x): # check that the predicted segmentation has improved self.assertGreater(diff_start, diff_end) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): target = torch.tensor([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]]) diff --git a/tests/test_generate_pos_neg_label_crop_centers.py b/tests/test_generate_pos_neg_label_crop_centers.py index e1d9398fe3..91db0e9d96 100644 --- a/tests/test_generate_pos_neg_label_crop_centers.py +++ b/tests/test_generate_pos_neg_label_crop_centers.py @@ -18,8 +18,7 @@ from monai.utils.misc import set_determinism from tests.utils import TEST_NDARRAYS, assert_allclose -TESTS = [] -TESTS.append( +TESTS = [ [ { "spatial_size": [2, 2, 2], @@ -33,7 +32,7 @@ 2, 3, ] -) +] class TestGeneratePosNegLabelCropCenters(unittest.TestCase): diff --git a/tests/test_get_equivalent_dtype.py b/tests/test_get_equivalent_dtype.py index fc0867523d..a4df3ac2ac 100644 --- a/tests/test_get_equivalent_dtype.py +++ b/tests/test_get_equivalent_dtype.py @@ -15,7 +15,7 @@ import torch from parameterized import parameterized -from monai.utils.type_conversion import get_equivalent_dtype +from monai.utils.type_conversion import get_equivalent_dtype, get_numpy_dtype_from_string, get_torch_dtype_from_string from tests.utils import TEST_NDARRAYS DTYPES = [torch.float32, np.float32, np.dtype(np.float32)] @@ -40,6 +40,16 @@ def test_native_type(self): out_dtype = get_equivalent_dtype(n, type(im_dtype)) self.assertEqual(out_dtype, n) + @parameterized.expand([["float", np.float64], ["float32", np.float32], ["float64", np.float64]]) + def test_from_string(self, dtype_str, expected_np): + expected_pt = get_equivalent_dtype(expected_np, torch.Tensor) + # numpy + dtype = get_numpy_dtype_from_string(dtype_str) + self.assertEqual(dtype, expected_np) + # torch + dtype = get_torch_dtype_from_string(dtype_str) + self.assertEqual(dtype, expected_pt) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_gibbs_noise.py b/tests/test_gibbs_noise.py index 3fbe047944..e40eda38db 100644 --- a/tests/test_gibbs_noise.py +++ b/tests/test_gibbs_noise.py @@ -13,14 +13,13 @@ from copy import deepcopy import numpy as np -import torch from parameterized import parameterized from monai.data.synthetic import create_test_image_2d, create_test_image_3d from monai.transforms import GibbsNoise from monai.utils.misc import set_determinism from monai.utils.module import optional_import -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose _, has_torch_fft = optional_import("torch.fft", name="fftshift") @@ -51,11 +50,7 @@ def test_same_result(self, im_shape, input_type): t = GibbsNoise(alpha) out1 = t(deepcopy(im)) out2 = t(deepcopy(im)) - self.assertEqual(type(out1), type(im)) - if isinstance(out1, torch.Tensor): - self.assertEqual(out1.device, im.device) - torch.testing.assert_allclose(out1, out2, rtol=1e-7, atol=0) - self.assertIsInstance(out1, type(im)) + assert_allclose(out1, out2, rtol=1e-7, atol=0, type_test="tensor") @parameterized.expand(TEST_CASES) def test_identity(self, im_shape, input_type): @@ -63,7 +58,7 @@ def test_identity(self, im_shape, input_type): alpha = 0.0 t = GibbsNoise(alpha) out = t(deepcopy(im)) - torch.testing.assert_allclose(im, out, atol=1e-2, rtol=1e-7) + assert_allclose(out, im, atol=1e-2, rtol=1e-7, type_test="tensor") @parameterized.expand(TEST_CASES) def test_alpha_1(self, im_shape, input_type): @@ -71,7 +66,7 @@ def test_alpha_1(self, im_shape, input_type): alpha = 1.0 t = GibbsNoise(alpha) out = t(deepcopy(im)) - torch.testing.assert_allclose(0 * im, out, rtol=1e-7, atol=0) + assert_allclose(out, 0 * im, rtol=1e-7, atol=0, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_gibbs_noised.py b/tests/test_gibbs_noised.py index 4905300703..6662e9e17c 100644 --- a/tests/test_gibbs_noised.py +++ b/tests/test_gibbs_noised.py @@ -13,14 +13,13 @@ from copy import deepcopy import numpy as np -import torch from parameterized import parameterized from monai.data.synthetic import create_test_image_2d, create_test_image_3d from monai.transforms import GibbsNoised from monai.utils.misc import set_determinism from monai.utils.module import optional_import -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose _, has_torch_fft = optional_import("torch.fft", name="fftshift") @@ -53,8 +52,7 @@ def test_same_result(self, im_shape, input_type): out1 = t(deepcopy(data)) out2 = t(deepcopy(data)) for k in KEYS: - torch.testing.assert_allclose(out1[k], out2[k], rtol=1e-7, atol=0) - self.assertIsInstance(out1[k], type(data[k])) + assert_allclose(out1[k], out2[k], rtol=1e-7, atol=0, type_test="tensor") @parameterized.expand(TEST_CASES) def test_identity(self, im_shape, input_type): @@ -63,11 +61,7 @@ def test_identity(self, im_shape, input_type): t = GibbsNoised(KEYS, alpha) out = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) - if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) - out[k], data[k] = out[k].cpu(), data[k].cpu() - np.testing.assert_allclose(data[k], out[k], atol=1e-2) + assert_allclose(out[k], data[k], atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_alpha_1(self, im_shape, input_type): @@ -76,11 +70,7 @@ def test_alpha_1(self, im_shape, input_type): t = GibbsNoised(KEYS, alpha) out = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) - if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) - out[k], data[k] = out[k].cpu(), data[k].cpu() - np.testing.assert_allclose(0.0 * data[k], out[k], atol=1e-2) + assert_allclose(out[k], 0.0 * data[k], atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_dict_matches(self, im_shape, input_type): @@ -89,7 +79,7 @@ def test_dict_matches(self, im_shape, input_type): alpha = 1.0 t = GibbsNoised(KEYS, alpha) out = t(deepcopy(data)) - torch.testing.assert_allclose(out[KEYS[0]], out[KEYS[1]], rtol=1e-7, atol=0) + assert_allclose(out[KEYS[0]], out[KEYS[1]], rtol=1e-7, atol=0, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_giou_loss.py b/tests/test_giou_loss.py new file mode 100644 index 0000000000..25cc258054 --- /dev/null +++ b/tests/test_giou_loss.py @@ -0,0 +1,59 @@ +# 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. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.losses import BoxGIoULoss + +TEST_CASES = [ + [ # shape: (1, 4), (1, 4) + {"input": torch.tensor([[1.0, 1.0, 2.0, 2.0]]), "target": torch.tensor([[1.0, 1.0, 2.0, 2.0]])}, + 0.0, + ], + [ # shape: (1, 6), (1, 6) + { + "input": torch.tensor([[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]]), + "target": torch.tensor([[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]]), + }, + 0.0, + ], +] + + +class TestGIoULoss(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_result(self, input_data, expected_val): + loss = BoxGIoULoss() + result = loss(**input_data) + np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4) + + def test_ill_shape(self): + loss = BoxGIoULoss() + with self.assertRaisesRegex(ValueError, ""): + loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3))) + + def test_with_cuda(self): + loss = BoxGIoULoss() + i = torch.tensor([[1.0, 1.0, 2.0, 2.0]]) + j = torch.tensor([[1.0, 1.0, 2.0, 2.0]]) + if torch.cuda.is_available(): + i = i.cuda() + j = j.cuda() + output = loss(i, j) + np.testing.assert_allclose(output.detach().cpu().numpy(), 0.0, atol=1e-4, rtol=1e-4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_global_mutual_information_loss.py b/tests/test_global_mutual_information_loss.py index b2a52f139b..73395d0572 100644 --- a/tests/test_global_mutual_information_loss.py +++ b/tests/test_global_mutual_information_loss.py @@ -18,7 +18,7 @@ from monai.losses.image_dissimilarity import GlobalMutualInformationLoss from tests.utils import SkipIfBeforePyTorchVersion, download_url_or_skip_test, skip_if_quick, testing_data_config -device = "cuda" if torch.cuda.is_available() else "cpu" +device = "cuda:0" if torch.cuda.is_available() else "cpu" FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", "temp_" + "mri.nii") diff --git a/tests/test_grid_dataset.py b/tests/test_grid_dataset.py index 9361d82cdf..30680c8e31 100644 --- a/tests/test_grid_dataset.py +++ b/tests/test_grid_dataset.py @@ -14,8 +14,8 @@ import numpy as np -from monai.data import DataLoader, GridPatchDataset, PatchIter -from monai.transforms import RandShiftIntensity +from monai.data import DataLoader, GridPatchDataset, PatchIter, PatchIterd +from monai.transforms import RandShiftIntensity, RandShiftIntensityd from monai.utils import set_determinism @@ -60,20 +60,62 @@ def test_loading_array(self): np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2)) np.testing.assert_allclose( item[0], - np.array([[[[1.4965, 2.4965], [5.4965, 6.4965]]], [[[11.3584, 12.3584], [15.3584, 16.3584]]]]), + np.array([[[[8.0577, 9.0577], [12.0577, 13.0577]]], [[[10.5540, 11.5540], [14.5540, 15.5540]]]]), rtol=1e-4, ) - np.testing.assert_allclose(item[1], np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5) + np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5) if sys.platform != "win32": for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=2): np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2)) np.testing.assert_allclose( item[0], - np.array([[[[1.2548, 2.2548], [5.2548, 6.2548]]], [[[9.1106, 10.1106], [13.1106, 14.1106]]]]), + np.array([[[[7.6533, 8.6533], [11.6533, 12.6533]]], [[[9.8524, 10.8524], [13.8524, 14.8524]]]]), rtol=1e-3, ) np.testing.assert_allclose( - item[1], np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5 + item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5 + ) + + def test_loading_dict(self): + set_determinism(seed=1234) + # test sequence input data with dict + data = [ + { + "image": np.arange(16, dtype=float).reshape(1, 4, 4), + "label": np.arange(16, dtype=float).reshape(1, 4, 4), + "metadata": "test string", + }, + { + "image": np.arange(16, dtype=float).reshape(1, 4, 4), + "label": np.arange(16, dtype=float).reshape(1, 4, 4), + "metadata": "test string", + }, + ] + # image level + patch_intensity = RandShiftIntensityd(keys="image", offsets=1.0, prob=1.0) + patch_iter = PatchIterd(keys=["image", "label"], patch_size=(2, 2), start_pos=(0, 0)) + ds = GridPatchDataset(data=data, patch_iter=patch_iter, transform=patch_intensity, with_coordinates=True) + # use the grid patch dataset + for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0): + np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2)) + np.testing.assert_equal(item[0]["label"].shape, (2, 1, 2, 2)) + self.assertListEqual(item[0]["metadata"], ["test string", "test string"]) + np.testing.assert_allclose( + item[0]["image"], + np.array([[[[8.0577, 9.0577], [12.0577, 13.0577]]], [[[10.5540, 11.5540], [14.5540, 15.5540]]]]), + rtol=1e-4, + ) + np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5) + if sys.platform != "win32": + for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=2): + np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2)) + np.testing.assert_allclose( + item[0]["image"], + np.array([[[[7.6533, 8.6533], [11.6533, 12.6533]]], [[[9.8524, 10.8524], [13.8524, 14.8524]]]]), + rtol=1e-3, + ) + np.testing.assert_allclose( + item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5 ) diff --git a/tests/test_grid_distortion.py b/tests/test_grid_distortion.py index 5e7ccd7c32..d71642aae8 100644 --- a/tests/test_grid_distortion.py +++ b/tests/test_grid_distortion.py @@ -15,10 +15,10 @@ from parameterized import parameterized from monai.transforms import GridDistortion -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TESTS.append( [ dict(num_cells=3, distort_steps=[(1.5,) * 4] * 2, mode="nearest", padding_mode="zeros"), @@ -101,7 +101,7 @@ class TestGridDistortion(unittest.TestCase): def test_grid_distortion(self, input_param, input_data, expected_val): g = GridDistortion(**input_param) result = g(input_data) - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4) + assert_allclose(result, expected_val, type_test=False, rtol=1e-4, atol=1e-4) if __name__ == "__main__": diff --git a/tests/test_grid_distortiond.py b/tests/test_grid_distortiond.py index 662596f935..2cf8bc7ff9 100644 --- a/tests/test_grid_distortiond.py +++ b/tests/test_grid_distortiond.py @@ -15,12 +15,12 @@ from parameterized import parameterized from monai.transforms import GridDistortiond -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] num_cells = (2, 2) distort_steps = [(1.5,) * (1 + n_c) for n_c in num_cells] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: img = np.indices([6, 6]).astype(np.float32) TESTS.append( [ @@ -77,8 +77,8 @@ class TestGridDistortiond(unittest.TestCase): def test_grid_distortiond(self, input_param, input_data, expected_val_img, expected_val_mask): g = GridDistortiond(**input_param) result = g(input_data) - assert_allclose(result["img"], expected_val_img, rtol=1e-4, atol=1e-4) - assert_allclose(result["mask"], expected_val_mask, rtol=1e-4, atol=1e-4) + assert_allclose(result["img"], expected_val_img, type_test=False, rtol=1e-4, atol=1e-4) + assert_allclose(result["mask"], expected_val_mask, type_test=False, rtol=1e-4, atol=1e-4) if __name__ == "__main__": diff --git a/tests/test_grid_patch.py b/tests/test_grid_patch.py new file mode 100644 index 0000000000..8a105afcd2 --- /dev/null +++ b/tests/test_grid_patch.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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.spatial.array import GridPatch +from tests.utils import TEST_NDARRAYS, assert_allclose + +A = np.arange(16).repeat(3).reshape(4, 4, 3).transpose(2, 0, 1) +A11 = A[:, :2, :2] +A12 = A[:, :2, 2:] +A21 = A[:, 2:, :2] +A22 = A[:, 2:, 2:] + +TEST_CASE_0 = [{"patch_size": (2, 2)}, A, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"patch_size": (2, 2), "num_patches": 3}, A, [A11, A12, A21]] +TEST_CASE_2 = [{"patch_size": (2, 2), "num_patches": 5}, A, [A11, A12, A21, A22, np.zeros((3, 2, 2))]] +TEST_CASE_3 = [{"patch_size": (2, 2), "offset": (0, 0)}, A, [A11, A12, A21, A22]] +TEST_CASE_4 = [{"patch_size": (2, 2), "offset": (0, 0)}, A, [A11, A12, A21, A22]] +TEST_CASE_5 = [{"patch_size": (2, 2), "offset": (2, 2)}, A, [A22]] +TEST_CASE_6 = [{"patch_size": (2, 2), "offset": (0, 2)}, A, [A12, A22]] +TEST_CASE_7 = [{"patch_size": (2, 2), "offset": (2, 0)}, A, [A21, A22]] +TEST_CASE_8 = [{"patch_size": (2, 2), "num_patches": 3, "sort_fn": "max"}, A, [A22, A21, A12]] +TEST_CASE_9 = [{"patch_size": (2, 2), "num_patches": 4, "sort_fn": "min"}, A, [A11, A12, A21, A22]] +TEST_CASE_10 = [{"patch_size": (2, 2), "overlap": 0.5, "num_patches": 3}, A, [A11, A[:, :2, 1:3], A12]] +TEST_CASE_11 = [ + {"patch_size": (3, 3), "num_patches": 2, "constant_values": 255}, + A, + [A[:, :3, :3], np.pad(A[:, :3, 3:], ((0, 0), (0, 0), (0, 2)), mode="constant", constant_values=255)], +] +TEST_CASE_12 = [ + {"patch_size": (3, 3), "offset": (-2, -2), "num_patches": 2}, + A, + [np.zeros((3, 3, 3)), np.pad(A[:, :1, 1:4], ((0, 0), (2, 0), (0, 0)), mode="constant")], +] +TEST_CASE_13 = [{"patch_size": (2, 2), "threshold": 50.0}, A, [A11]] + + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + TEST_SINGLE.append([p, *TEST_CASE_9]) + TEST_SINGLE.append([p, *TEST_CASE_10]) + TEST_SINGLE.append([p, *TEST_CASE_11]) + TEST_SINGLE.append([p, *TEST_CASE_12]) + TEST_SINGLE.append([p, *TEST_CASE_13]) + + +class TestGridPatch(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_grid_patch(self, in_type, input_parameters, image, expected): + input_image = in_type(image) + splitter = GridPatch(**input_parameters) + output = list(splitter(input_image)) + self.assertEqual(len(output), len(expected)) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[0], expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_grid_patchd.py b/tests/test_grid_patchd.py new file mode 100644 index 0000000000..8f1e238b42 --- /dev/null +++ b/tests/test_grid_patchd.py @@ -0,0 +1,84 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.spatial.dictionary import GridPatchd +from tests.utils import TEST_NDARRAYS, assert_allclose + +A = np.arange(16).repeat(3).reshape(4, 4, 3).transpose(2, 0, 1) +A11 = A[:, :2, :2] +A12 = A[:, :2, 2:] +A21 = A[:, 2:, :2] +A22 = A[:, 2:, 2:] + +TEST_CASE_0 = [{"patch_size": (2, 2)}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"patch_size": (2, 2), "num_patches": 3}, {"image": A}, [A11, A12, A21]] +TEST_CASE_2 = [{"patch_size": (2, 2), "num_patches": 5}, {"image": A}, [A11, A12, A21, A22, np.zeros((3, 2, 2))]] +TEST_CASE_3 = [{"patch_size": (2, 2), "offset": (0, 0)}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_4 = [{"patch_size": (2, 2), "offset": (0, 0)}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_5 = [{"patch_size": (2, 2), "offset": (2, 2)}, {"image": A}, [A22]] +TEST_CASE_6 = [{"patch_size": (2, 2), "offset": (0, 2)}, {"image": A}, [A12, A22]] +TEST_CASE_7 = [{"patch_size": (2, 2), "offset": (2, 0)}, {"image": A}, [A21, A22]] +TEST_CASE_8 = [{"patch_size": (2, 2), "num_patches": 3, "sort_fn": "max"}, {"image": A}, [A22, A21, A12]] +TEST_CASE_9 = [{"patch_size": (2, 2), "num_patches": 4, "sort_fn": "min"}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_10 = [{"patch_size": (2, 2), "overlap": 0.5, "num_patches": 3}, {"image": A}, [A11, A[:, :2, 1:3], A12]] +TEST_CASE_11 = [ + {"patch_size": (3, 3), "num_patches": 2, "constant_values": 255}, + {"image": A}, + [A[:, :3, :3], np.pad(A[:, :3, 3:], ((0, 0), (0, 0), (0, 2)), mode="constant", constant_values=255)], +] +TEST_CASE_12 = [ + {"patch_size": (3, 3), "offset": (-2, -2), "num_patches": 2}, + {"image": A}, + [np.zeros((3, 3, 3)), np.pad(A[:, :1, 1:4], ((0, 0), (2, 0), (0, 0)), mode="constant")], +] +TEST_CASE_13 = [{"patch_size": (2, 2), "threshold": 50.0}, {"image": A}, [A11]] + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + TEST_SINGLE.append([p, *TEST_CASE_9]) + TEST_SINGLE.append([p, *TEST_CASE_10]) + TEST_SINGLE.append([p, *TEST_CASE_11]) + TEST_SINGLE.append([p, *TEST_CASE_12]) + TEST_SINGLE.append([p, *TEST_CASE_13]) + + +class TestGridPatchd(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_grid_patchd(self, in_type, input_parameters, image_dict, expected): + image_key = "image" + input_dict = {} + for k, v in image_dict.items(): + input_dict[k] = v + if k == image_key: + input_dict[k] = in_type(v) + splitter = GridPatchd(keys=image_key, **input_parameters) + output = list(splitter(input_dict)) + self.assertEqual(len(output), len(expected)) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[image_key], expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_grid_split.py b/tests/test_grid_split.py new file mode 100644 index 0000000000..82734ffd93 --- /dev/null +++ b/tests/test_grid_split.py @@ -0,0 +1,86 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.transforms import GridSplit +from tests.utils import TEST_NDARRAYS, assert_allclose + +A11 = torch.randn(3, 2, 2) +A12 = torch.randn(3, 2, 2) +A21 = torch.randn(3, 2, 2) +A22 = torch.randn(3, 2, 2) + +A1 = torch.cat([A11, A12], 2) +A2 = torch.cat([A21, A22], 2) +A = torch.cat([A1, A2], 1) + +TEST_CASE_0 = [{"grid": (2, 2)}, A, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"grid": (2, 1)}, A, [A1, A2]] +TEST_CASE_2 = [{"grid": (1, 2)}, A1, [A11, A12]] +TEST_CASE_3 = [{"grid": (1, 2)}, A2, [A21, A22]] +TEST_CASE_4 = [{"grid": (1, 1), "size": (2, 2)}, A, [A11]] +TEST_CASE_5 = [{"grid": (1, 1), "size": 4}, A, [A]] +TEST_CASE_6 = [{"grid": (2, 2), "size": 2}, A, [A11, A12, A21, A22]] +TEST_CASE_7 = [{"grid": (1, 1)}, A, [A]] +TEST_CASE_8 = [ + {"grid": (2, 2), "size": 2}, + torch.arange(12).reshape(1, 3, 4).to(torch.float32), + torch.Tensor([[[[0, 1], [4, 5]]], [[[2, 3], [6, 7]]], [[[4, 5], [8, 9]]], [[[6, 7], [10, 11]]]]).to(torch.float32), +] + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + +TEST_CASE_MC_0 = [{"grid": (2, 2)}, [A, A], [[A11, A12, A21, A22], [A11, A12, A21, A22]]] +TEST_CASE_MC_1 = [{"grid": (2, 1)}, [A] * 5, [[A1, A2]] * 5] +TEST_CASE_MC_2 = [{"grid": (1, 2)}, [A1, A2], [[A11, A12], [A21, A22]]] + +TEST_MULTIPLE = [] +for p in TEST_NDARRAYS: + TEST_MULTIPLE.append([p, *TEST_CASE_MC_0]) + TEST_MULTIPLE.append([p, *TEST_CASE_MC_1]) + TEST_MULTIPLE.append([p, *TEST_CASE_MC_2]) + + +class TestGridSplit(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_split_patch_single_call(self, in_type, input_parameters, image, expected): + input_image = in_type(image) + splitter = GridSplit(**input_parameters) + output = splitter(input_image) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch, expected_patch, type_test=False) + + @parameterized.expand(TEST_MULTIPLE) + def test_split_patch_multiple_call(self, in_type, input_parameters, img_list, expected_list): + splitter = GridSplit(**input_parameters) + for image, expected in zip(img_list, expected_list): + input_image = in_type(image) + output = splitter(input_image) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch, expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_grid_splitd.py b/tests/test_grid_splitd.py new file mode 100644 index 0000000000..086dd2691d --- /dev/null +++ b/tests/test_grid_splitd.py @@ -0,0 +1,94 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.transforms import GridSplitd +from tests.utils import TEST_NDARRAYS, assert_allclose + +A11 = torch.randn(3, 2, 2) +A12 = torch.randn(3, 2, 2) +A21 = torch.randn(3, 2, 2) +A22 = torch.randn(3, 2, 2) + +A1 = torch.cat([A11, A12], 2) +A2 = torch.cat([A21, A22], 2) +A = torch.cat([A1, A2], 1) + +TEST_CASE_0 = [{"keys": "image", "grid": (2, 2)}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"keys": "image", "grid": (2, 1)}, {"image": A}, [A1, A2]] +TEST_CASE_2 = [{"keys": "image", "grid": (1, 2)}, {"image": A1}, [A11, A12]] +TEST_CASE_3 = [{"keys": "image", "grid": (1, 2)}, {"image": A2}, [A21, A22]] +TEST_CASE_4 = [{"keys": "image", "grid": (1, 1), "size": {"image": (2, 2)}}, {"image": A}, [A11]] +TEST_CASE_5 = [{"keys": "image", "grid": (1, 1), "size": {"image": 4}}, {"image": A}, [A]] +TEST_CASE_6 = [{"keys": "image", "grid": (2, 2), "size": {"image": 2}}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_7 = [{"keys": "image", "grid": (1, 1)}, {"image": A}, [A]] +TEST_CASE_8 = [ + {"keys": "image", "grid": (2, 2), "size": {"image": 2}}, + {"image": torch.arange(12).reshape(1, 3, 4).to(torch.float32)}, + torch.Tensor([[[[0, 1], [4, 5]]], [[[2, 3], [6, 7]]], [[[4, 5], [8, 9]]], [[[6, 7], [10, 11]]]]).to(torch.float32), +] + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + +TEST_CASE_MC_0 = [ + {"keys": "image", "grid": (2, 2)}, + [{"image": A}, {"image": A}], + [[A11, A12, A21, A22], [A11, A12, A21, A22]], +] +TEST_CASE_MC_1 = [{"keys": "image", "grid": (2, 1)}, [{"image": A}, {"image": A}, {"image": A}], [[A1, A2]] * 3] +TEST_CASE_MC_2 = [{"keys": "image", "grid": (1, 2)}, [{"image": A1}, {"image": A2}], [[A11, A12], [A21, A22]]] + +TEST_MULTIPLE = [] +for p in TEST_NDARRAYS: + TEST_MULTIPLE.append([p, *TEST_CASE_MC_0]) + TEST_MULTIPLE.append([p, *TEST_CASE_MC_1]) + TEST_MULTIPLE.append([p, *TEST_CASE_MC_2]) + + +class TestGridSplitd(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_split_patch_single_call(self, in_type, input_parameters, img_dict, expected): + input_dict = {} + for k, v in img_dict.items(): + input_dict[k] = in_type(v) + splitter = GridSplitd(**input_parameters) + output = splitter(input_dict) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[input_parameters["keys"]], expected_patch, type_test=False) + + @parameterized.expand(TEST_MULTIPLE) + def test_split_patch_multiple_call(self, in_type, input_parameters, img_list, expected_list): + splitter = GridSplitd(**input_parameters) + for img_dict, expected in zip(img_list, expected_list): + input_dict = {} + for k, v in img_dict.items(): + input_dict[k] = in_type(v) + output = splitter(input_dict) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[input_parameters["keys"]], expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_handler_prob_map_producer.py b/tests/test_handler_prob_map_producer.py index b3f79cf587..399657ac60 100644 --- a/tests/test_handler_prob_map_producer.py +++ b/tests/test_handler_prob_map_producer.py @@ -18,30 +18,48 @@ from parameterized import parameterized from torch.utils.data import DataLoader -from monai.apps.pathology.handlers import ProbMapProducer from monai.data.dataset import Dataset from monai.engines import Evaluator -from monai.handlers import ValidationHandler +from monai.handlers import ProbMapProducer, ValidationHandler +from monai.utils.enums import ProbMapKeys TEST_CASE_0 = ["temp_image_inference_output_1", 2] TEST_CASE_1 = ["temp_image_inference_output_2", 9] -TEST_CASE_2 = ["temp_image_inference_output_3", 1000] +TEST_CASE_2 = ["temp_image_inference_output_3", 100] class TestDataset(Dataset): def __init__(self, name, size): super().__init__( data=[ - {"name": name, "mask_shape": (size, size), "mask_locations": [[i, i] for i in range(size)], "level": 0} + { + "image": name, + ProbMapKeys.COUNT.value: size, + ProbMapKeys.SIZE.value: np.array([size, size]), + ProbMapKeys.LOCATION.value: np.array([i, i]), + } + for i in range(size) ] ) - self.len = size - - def __len__(self): - return self.len + self.image_data = [ + { + ProbMapKeys.NAME.value: name, + ProbMapKeys.COUNT.value: size, + ProbMapKeys.SIZE.value: np.array([size, size]), + } + ] def __getitem__(self, index): - return {"name": self.data[0]["name"], "mask_location": self.data[0]["mask_locations"][index], "pred": index + 1} + return { + "image": np.zeros((3, 2, 2)), + ProbMapKeys.COUNT.value: self.data[index][ProbMapKeys.COUNT.value], + "metadata": { + ProbMapKeys.NAME.value: self.data[index]["image"], + ProbMapKeys.SIZE.value: self.data[index][ProbMapKeys.SIZE.value], + ProbMapKeys.LOCATION.value: self.data[index][ProbMapKeys.LOCATION.value], + }, + "pred": index + 1, + } class TestEvaluator(Evaluator): diff --git a/tests/test_handler_segmentation_saver.py b/tests/test_handler_segmentation_saver.py deleted file mode 100644 index ee6566f6cb..0000000000 --- a/tests/test_handler_segmentation_saver.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. - -import os -import tempfile -import unittest - -import numpy as np -import torch -from ignite.engine import Engine -from parameterized import parameterized - -from monai.data import decollate_batch -from monai.handlers import SegmentationSaver - -TEST_CASE_0 = [".nii.gz"] - -TEST_CASE_1 = [".png"] - - -class TestHandlerSegmentationSaver(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1]) - def test_saved_content(self, output_ext): - with tempfile.TemporaryDirectory() as tempdir: - - # set up engine - def _train_func(engine, batch): - engine.state.batch = decollate_batch(batch) - return [torch.randint(0, 255, (1, 2, 2)).float() for _ in range(8)] - - engine = Engine(_train_func) - - # set up testing handler - saver = SegmentationSaver( - output_dir=tempdir, output_postfix="seg", output_ext=output_ext, scale=255, output_dtype=np.uint8 - ) - saver.attach(engine) - - data = [ - { - "filename_or_obj": ["testfile" + str(i) + ".nii.gz" for i in range(8)], - "patch_index": torch.tensor(list(range(8))), - } - ] - engine.run(data, max_epochs=1) - for i in range(8): - filepath = os.path.join("testfile" + str(i), "testfile" + str(i) + "_seg" + f"_{i}" + output_ext) - self.assertTrue(os.path.exists(os.path.join(tempdir, filepath))) - - @parameterized.expand([TEST_CASE_0, TEST_CASE_1]) - def test_save_resized_content(self, output_ext): - with tempfile.TemporaryDirectory() as tempdir: - - # set up engine - def _train_func(engine, batch): - engine.state.batch = decollate_batch(batch) - return [torch.randint(0, 255, (1, 2, 2)).float() for _ in range(8)] - - engine = Engine(_train_func) - - # set up testing handler - saver = SegmentationSaver( - output_dir=tempdir, output_postfix="seg", output_ext=output_ext, scale=255, output_dtype=np.uint8 - ) - saver.attach(engine) - - data = [ - { - "filename_or_obj": ["testfile" + str(i) + ".nii.gz" for i in range(8)], - "spatial_shape": torch.tensor([[28, 28] for _ in range(8)]), - "affine": torch.tensor([np.diag(np.ones(4)) * 5 for _ in range(8)]), - "original_affine": torch.tensor([np.diag(np.ones(4)) * 1.0 for _ in range(8)]), - } - ] - engine.run(data, max_epochs=1) - for i in range(8): - filepath = os.path.join("testfile" + str(i), "testfile" + str(i) + "_seg" + output_ext) - self.assertTrue(os.path.exists(os.path.join(tempdir, filepath))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_handler_smartcache.py b/tests/test_handler_smartcache.py index ec96d47e3d..7bc9011c2d 100644 --- a/tests/test_handler_smartcache.py +++ b/tests/test_handler_smartcache.py @@ -17,10 +17,8 @@ from monai.data import SmartCacheDataset from monai.handlers import SmartCacheHandler -from tests.utils import SkipIfBeforePyTorchVersion -@SkipIfBeforePyTorchVersion((1, 7)) class TestHandlerSmartCache(unittest.TestCase): def test_content(self): data = [0, 1, 2, 3, 4, 5, 6, 7, 8] diff --git a/tests/test_handler_tb_image.py b/tests/test_handler_tb_image.py index d11bbfec59..749480e279 100644 --- a/tests/test_handler_tb_image.py +++ b/tests/test_handler_tb_image.py @@ -20,10 +20,14 @@ from monai.data import decollate_batch from monai.handlers import TensorBoardImageHandler +from monai.utils import optional_import + +_, has_tb = optional_import("torch.utils.tensorboard", name="SummaryWriter") TEST_CASES = [[[20, 20]], [[2, 20, 20]], [[3, 20, 20]], [[20, 20, 20]], [[2, 20, 20, 20]], [[2, 2, 20, 20, 20]]] +@unittest.skipUnless(has_tb, "Requires SummaryWriter installation") class TestHandlerTBImage(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_tb_image_shape(self, shape): diff --git a/tests/test_handler_tb_stats.py b/tests/test_handler_tb_stats.py index 4d582d151b..eef77e5e2b 100644 --- a/tests/test_handler_tb_stats.py +++ b/tests/test_handler_tb_stats.py @@ -14,11 +14,14 @@ import unittest from ignite.engine import Engine, Events -from torch.utils.tensorboard import SummaryWriter from monai.handlers import TensorBoardStatsHandler +from monai.utils import optional_import +SummaryWriter, has_tb = optional_import("torch.utils.tensorboard", name="SummaryWriter") + +@unittest.skipUnless(has_tb, "Requires SummaryWriter installation") class TestHandlerTBStats(unittest.TestCase): def test_metrics_print(self): with tempfile.TemporaryDirectory() as tempdir: diff --git a/tests/test_hardnegsampler.py b/tests/test_hardnegsampler.py new file mode 100644 index 0000000000..f4eff81810 --- /dev/null +++ b/tests/test_hardnegsampler.py @@ -0,0 +1,52 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.hard_negative_sampler import HardNegativeSampler +from tests.utils import assert_allclose + +TEST_CASE = [ + [[], [], [], [torch.tensor([]), torch.tensor([])], [torch.tensor([]), torch.tensor([])]], + [ + [0, 1], + [1, 0, 2, 3], + [0.1, 0.9, 0.4, 0.3, 0.3, 0.5], + [torch.tensor([0, 1]), torch.tensor([1, 0, 1, 1])], + [torch.tensor([1, 0]), torch.tensor([0, 1, 0, 0])], + ], +] + +select_sample_size_per_image = 6 +positive_fraction = 0.5 +min_neg = 1 +pool_size = 2 + + +class TestSampleSlices(unittest.TestCase): + @parameterized.expand(TEST_CASE) + def test_shape(self, target_label0, target_label1, concat_fg_probs, expected_result_pos, expected_result_neg): + compute_dtypes = [torch.float16, torch.float32] + for compute_dtype in compute_dtypes: + sampler = HardNegativeSampler(select_sample_size_per_image, positive_fraction, min_neg, pool_size) + target_labels = [torch.tensor(target_label0), torch.tensor(target_label1)] + result_pos, result_neg = sampler(target_labels, torch.tensor(concat_fg_probs, dtype=compute_dtype)) + for r, er in zip(result_pos, expected_result_pos): + assert_allclose(r, er) + for r, er in zip(result_neg, expected_result_neg): + assert_allclose(r, er) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_hausdorff_distance.py b/tests/test_hausdorff_distance.py index 79a2c84b37..44c011fe13 100644 --- a/tests/test_hausdorff_distance.py +++ b/tests/test_hausdorff_distance.py @@ -127,7 +127,7 @@ def test_value(self, input_data, expected_value): batch_seg_1 = seg_1.unsqueeze(0).unsqueeze(0).repeat([batch, n_class, 1, 1, 1]) batch_seg_2 = seg_2.unsqueeze(0).unsqueeze(0).repeat([batch, n_class, 1, 1, 1]) hd_metric(batch_seg_1, batch_seg_2) - result = hd_metric.aggregate() + result = hd_metric.aggregate(reduction="mean") expected_value_curr = expected_value[ct] np.testing.assert_allclose(expected_value_curr, result, rtol=1e-7) ct += 1 diff --git a/tests/test_hilbert_transform.py b/tests/test_hilbert_transform.py index 10aa83293f..f7954d6b24 100644 --- a/tests/test_hilbert_transform.py +++ b/tests/test_hilbert_transform.py @@ -16,14 +16,8 @@ from parameterized import parameterized from monai.networks.layers import HilbertTransform -from monai.utils import InvalidPyTorchVersionError, OptionalImportError -from tests.utils import ( - SkipIfAtLeastPyTorchVersion, - SkipIfBeforePyTorchVersion, - SkipIfModule, - SkipIfNoModule, - skip_if_no_cuda, -) +from monai.utils import OptionalImportError +from tests.utils import SkipIfModule, SkipIfNoModule, skip_if_no_cuda def create_expected_numpy_output(input_datum, **kwargs): @@ -164,7 +158,6 @@ def create_expected_numpy_output(input_datum, **kwargs): # TESTS CHECKING PADDING, AXIS SELECTION ETC ARE COVERED BY test_detect_envelope.py -@SkipIfBeforePyTorchVersion((1, 7)) @SkipIfNoModule("torch.fft") class TestHilbertTransformCPU(unittest.TestCase): @parameterized.expand( @@ -183,7 +176,6 @@ def test_value(self, arguments, image, expected_data, atol): @skip_if_no_cuda -@SkipIfBeforePyTorchVersion((1, 7)) @SkipIfNoModule("torch.fft") class TestHilbertTransformGPU(unittest.TestCase): @parameterized.expand( @@ -204,19 +196,11 @@ def test_value(self, arguments, image, expected_data, atol): np.testing.assert_allclose(result, expected_data.squeeze(), atol=atol) -@SkipIfBeforePyTorchVersion((1, 7)) @SkipIfModule("torch.fft") class TestHilbertTransformNoFFTMod(unittest.TestCase): def test_no_fft_module_error(self): self.assertRaises(OptionalImportError, HilbertTransform(), torch.randn(1, 1, 10)) -@SkipIfAtLeastPyTorchVersion((1, 7)) -class TestHilbertTransformInvalidPyTorch(unittest.TestCase): - def test_invalid_pytorch_error(self): - with self.assertRaisesRegex(InvalidPyTorchVersionError, "version"): - HilbertTransform() - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_histogram_normalize.py b/tests/test_histogram_normalize.py index 95aa37f26e..9218d247b1 100644 --- a/tests/test_histogram_normalize.py +++ b/tests/test_histogram_normalize.py @@ -49,7 +49,7 @@ class TestHistogramNormalize(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = HistogramNormalize(**argments)(image) - assert_allclose(result, expected_data) + assert_allclose(result, expected_data, type_test="tensor") self.assertEqual(get_equivalent_dtype(result.dtype, data_type=np.ndarray), argments.get("dtype", np.float32)) diff --git a/tests/test_histogram_normalized.py b/tests/test_histogram_normalized.py index 7b86a9685f..a56b063847 100644 --- a/tests/test_histogram_normalized.py +++ b/tests/test_histogram_normalized.py @@ -49,7 +49,7 @@ class TestHistogramNormalized(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = HistogramNormalized(**argments)(image)["img"] - assert_allclose(result, expected_data) + assert_allclose(result, expected_data, type_test="tensor") self.assertEqual(get_equivalent_dtype(result.dtype, data_type=np.ndarray), argments.get("dtype", np.float32)) diff --git a/tests/test_hovernet.py b/tests/test_hovernet.py new file mode 100644 index 0000000000..568aeb04dc --- /dev/null +++ b/tests/test_hovernet.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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks import eval_mode +from monai.networks.nets import HoverNet +from tests.utils import test_script_save + +device = "cuda" if torch.cuda.is_available() else "cpu" + +TEST_CASE_0 = [ # fast mode, batch 16 + {"out_classes": 5, "mode": HoverNet.Mode.FAST}, + (1, 3, 256, 256), + { + "nucleus_prediction": (1, 2, 164, 164), + "type_prediction": (1, 5, 164, 164), + "horizonal_vertical": (1, 2, 164, 164), + }, +] + +TEST_CASE_1 = [ # single channel 2D, batch 16 + {"mode": HoverNet.Mode.FAST}, + (1, 3, 256, 256), + {"nucleus_prediction": (1, 2, 164, 164), "horizonal_vertical": (1, 2, 164, 164)}, +] + +TEST_CASE_2 = [ # single channel 3D, batch 16 + {"mode": HoverNet.Mode.ORIGINAL}, + (1, 3, 270, 270), + {"nucleus_prediction": (1, 2, 80, 80), "horizonal_vertical": (1, 2, 80, 80)}, +] + +TEST_CASE_3 = [ # 4-channel 3D, batch 16 + {"out_classes": 6, "mode": HoverNet.Mode.ORIGINAL}, + (1, 3, 270, 270), + {"nucleus_prediction": (1, 2, 80, 80), "type_prediction": (1, 6, 80, 80), "horizonal_vertical": (1, 2, 80, 80)}, +] + +TEST_CASE_4 = [ # 4-channel 3D, batch 16, batch normalization + {"mode": HoverNet.Mode.FAST, "dropout_prob": 0.5}, + (1, 3, 256, 256), + {"nucleus_prediction": (1, 2, 164, 164), "horizonal_vertical": (1, 2, 164, 164)}, +] + +CASES = [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4] + +ILL_CASES = [ + [{"out_classes": 6, "mode": 3}], + [{"out_classes": 1000, "mode": HoverNet.Mode.ORIGINAL}], + [{"out_classes": 1, "mode": HoverNet.Mode.ORIGINAL}], + [{"out_classes": 6, "mode": HoverNet.Mode.ORIGINAL, "dropout_prob": 100}], +] + + +class TestHoverNet(unittest.TestCase): + @parameterized.expand(CASES) + def test_shape(self, input_param, input_shape, expected_shapes): + net = HoverNet(**input_param).to(device) + with eval_mode(net): + result = net.forward(torch.randn(input_shape).to(device)) + for item in result: + self.assertEqual(result[item].shape, expected_shapes[item]) + + def test_script(self): + net = HoverNet(mode=HoverNet.Mode.FAST) + test_data = torch.randn(1, 3, 256, 256) + test_script_save(net, test_data) + + def test_script_without_running_stats(self): + net = HoverNet(mode=HoverNet.Mode.FAST) + test_data = torch.randn(1, 3, 256, 256) + test_script_save(net, test_data) + + def test_ill_input_shape(self): + net = HoverNet(mode=HoverNet.Mode.FAST) + with eval_mode(net): + with self.assertRaises(ValueError): + net.forward(torch.randn(1, 3, 270, 260)) + + @parameterized.expand(ILL_CASES) + def test_ill_input_hyper_params(self, input_param): + with self.assertRaises(ValueError): + _ = HoverNet(**input_param) + + +if __name__ == "__main__": + unittest.main(argv=["first-arg-is-ignored"], exit=False) diff --git a/tests/test_identityd.py b/tests/test_identityd.py index 2df74ba2c6..f1d27d61d4 100644 --- a/tests/test_identityd.py +++ b/tests/test_identityd.py @@ -19,8 +19,7 @@ class TestIdentityd(NumpyImageTestCase2D): def test_identityd(self): for p in TEST_NDARRAYS: img = p(self.imt) - data = {} - data["img"] = img + data = {"img": img} identity = Identityd(keys=data.keys()) assert_allclose(img, identity(data)["img"]) diff --git a/tests/test_image_dataset.py b/tests/test_image_dataset.py index 41eda803dc..a89759323d 100644 --- a/tests/test_image_dataset.py +++ b/tests/test_image_dataset.py @@ -15,6 +15,7 @@ import nibabel as nib import numpy as np +import torch from monai.data import ImageDataset from monai.transforms import ( @@ -45,8 +46,9 @@ def __call__(self, data): class _TestCompose(Compose): def __call__(self, data, meta): - data = self.transforms[0](data, meta) # ensure channel first - data, _, meta["affine"] = self.transforms[1](data, meta["affine"]) # spacing + data = self.transforms[0](data) # ensure channel first + data = self.transforms[1](data, data.meta["affine"]) # spacing + meta = data.meta if len(self.transforms) == 3: return self.transforms[2](data), meta # image contrast return data, meta @@ -55,8 +57,8 @@ def __call__(self, data, meta): class TestImageDataset(unittest.TestCase): def test_use_case(self): with tempfile.TemporaryDirectory() as tempdir: - img_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)), np.eye(4)) - seg_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)), np.eye(4)) + img_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)).astype(float), np.eye(4)) + seg_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)).astype(float), np.eye(4)) img_name, seg_name = os.path.join(tempdir, "img.nii.gz"), os.path.join(tempdir, "seg.nii.gz") nib.save(img_, img_name) nib.save(seg_, seg_name) @@ -79,7 +81,7 @@ def test_dataset(self): with tempfile.TemporaryDirectory() as tempdir: full_names, ref_data = [], [] for filename in FILENAMES: - test_image = np.random.randint(0, 2, size=(4, 4, 4)) + test_image = np.random.randint(0, 2, size=(4, 4, 4)).astype(float) ref_data.append(test_image) save_path = os.path.join(tempdir, filename) full_names.append(save_path) @@ -93,7 +95,7 @@ def test_dataset(self): # loading no meta, int dataset = ImageDataset(full_names, dtype=np.float16) for d, _ in zip(dataset, ref_data): - self.assertEqual(d.dtype, np.float16) + self.assertEqual(d.dtype, torch.float16) # loading with meta, no transform dataset = ImageDataset(full_names, image_only=False) diff --git a/tests/test_image_rw.py b/tests/test_image_rw.py index 62b1147aa5..f234b32c24 100644 --- a/tests/test_image_rw.py +++ b/tests/test_image_rw.py @@ -16,10 +16,12 @@ import unittest import numpy as np +import torch from parameterized import parameterized -from monai.data.image_reader import ITKReader, NibabelReader, PILReader +from monai.data.image_reader import ITKReader, NibabelReader, NrrdReader, PILReader from monai.data.image_writer import ITKWriter, NibabelWriter, PILWriter, register_writer, resolve_writer +from monai.data.meta_tensor import MetaTensor from monai.transforms import LoadImage, SaveImage, moveaxis from monai.utils import OptionalImportError from tests.utils import TEST_NDARRAYS, assert_allclose @@ -41,25 +43,25 @@ def nifti_rw(self, test_data, reader, writer, dtype, resample=True): saver = SaveImage( output_dir=self.test_dir, output_ext=output_ext, resample=resample, separate_folder=False, writer=writer ) - saver( - p(test_data), - { - "filename_or_obj": f"{filepath}.png", - "affine": np.eye(4), - "original_affine": np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]), - }, - ) + meta_dict = { + "filename_or_obj": f"{filepath}.png", + "affine": np.eye(4), + "original_affine": np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]), + } + test_data = MetaTensor(p(test_data), meta=meta_dict) + saver(test_data) saved_path = os.path.join(self.test_dir, filepath + "_trans" + output_ext) self.assertTrue(os.path.exists(saved_path)) - loader = LoadImage(reader=reader, squeeze_non_spatial_dims=True) - data, meta = loader(saved_path) + loader = LoadImage(image_only=True, reader=reader, squeeze_non_spatial_dims=True) + data = loader(saved_path) + meta = data.meta if meta["original_channel_dim"] == -1: _test_data = moveaxis(test_data, 0, -1) else: _test_data = test_data[0] if resample: _test_data = moveaxis(_test_data, 0, 1) - assert_allclose(data, _test_data) + assert_allclose(data, torch.as_tensor(_test_data)) @parameterized.expand(itertools.product([NibabelReader, ITKReader], [NibabelWriter, "ITKWriter"])) def test_2d(self, reader, writer): @@ -95,16 +97,18 @@ def png_rw(self, test_data, reader, writer, dtype, resample=True): saver = SaveImage( output_dir=self.test_dir, output_ext=output_ext, resample=resample, separate_folder=False, writer=writer ) - saver(p(test_data), {"filename_or_obj": f"{filepath}.png", "spatial_shape": (6, 8)}) + test_data = MetaTensor(p(test_data), meta={"filename_or_obj": f"{filepath}.png", "spatial_shape": (6, 8)}) + saver(test_data) saved_path = os.path.join(self.test_dir, filepath + "_trans" + output_ext) self.assertTrue(os.path.exists(saved_path)) - loader = LoadImage(reader=reader) - data, meta = loader(saved_path) + loader = LoadImage(image_only=True, reader=reader) + data = loader(saved_path) + meta = data.meta if meta["original_channel_dim"] == -1: _test_data = moveaxis(test_data, 0, -1) else: _test_data = test_data[0] - assert_allclose(data, _test_data) + assert_allclose(data, torch.as_tensor(_test_data)) @parameterized.expand(itertools.product([PILReader, ITKReader], [PILWriter, ITKWriter])) def test_2d(self, reader, writer): @@ -132,5 +136,41 @@ def test_1_new(self): self.assertEqual(resolve_writer("new")[0](0), 1) +class TestLoadSaveNrrd(unittest.TestCase): + def setUp(self): + self.test_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def nrrd_rw(self, test_data, reader, writer, dtype, resample=True): + test_data = test_data.astype(dtype) + ndim = len(test_data.shape) + for p in TEST_NDARRAYS: + output_ext = ".nrrd" + filepath = f"testfile_{ndim}d" + saver = SaveImage( + output_dir=self.test_dir, output_ext=output_ext, resample=resample, separate_folder=False, writer=writer + ) + test_data = MetaTensor( + p(test_data), meta={"filename_or_obj": f"{filepath}{output_ext}", "spatial_shape": test_data.shape} + ) + saver(test_data) + saved_path = os.path.join(self.test_dir, filepath + "_trans" + output_ext) + loader = LoadImage(image_only=True, reader=reader) + data = loader(saved_path) + assert_allclose(data, torch.as_tensor(test_data)) + + @parameterized.expand(itertools.product([NrrdReader, ITKReader], [ITKWriter, ITKWriter])) + def test_2d(self, reader, writer): + test_data = np.random.randn(8, 8).astype(np.float32) + self.nrrd_rw(test_data, reader, writer, np.float32) + + @parameterized.expand(itertools.product([NrrdReader, ITKReader], [ITKWriter, ITKWriter])) + def test_3d(self, reader, writer): + test_data = np.random.randn(8, 8, 8).astype(np.float32) + self.nrrd_rw(test_data, reader, writer, np.float32) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_init_reader.py b/tests/test_init_reader.py index 03a63cc375..e9a8ec96f7 100644 --- a/tests/test_init_reader.py +++ b/tests/test_init_reader.py @@ -11,7 +11,7 @@ import unittest -from monai.data import ITKReader, NibabelReader, NumpyReader, PILReader +from monai.data import ITKReader, NibabelReader, NrrdReader, NumpyReader, PILReader, PydicomReader from monai.transforms import LoadImage, LoadImaged from tests.utils import SkipIfNoModule @@ -23,13 +23,15 @@ def test_load_image(self): self.assertIsInstance(instance1, LoadImage) self.assertIsInstance(instance2, LoadImage) - for r in ["NibabelReader", "PILReader", "ITKReader", "NumpyReader", None]: + for r in ["NibabelReader", "PILReader", "ITKReader", "NumpyReader", "NrrdReader", "PydicomReader", None]: inst = LoadImaged("image", reader=r) self.assertIsInstance(inst, LoadImaged) @SkipIfNoModule("itk") @SkipIfNoModule("nibabel") @SkipIfNoModule("PIL") + @SkipIfNoModule("nrrd") + @SkipIfNoModule("Pydicom") def test_readers(self): inst = ITKReader() self.assertIsInstance(inst, ITKReader) @@ -39,6 +41,9 @@ def test_readers(self): inst = NibabelReader(as_closest_canonical=True) self.assertIsInstance(inst, NibabelReader) + inst = PydicomReader() + self.assertIsInstance(inst, PydicomReader) + inst = NumpyReader() self.assertIsInstance(inst, NumpyReader) inst = NumpyReader(npz_keys="test") @@ -47,6 +52,9 @@ def test_readers(self): inst = PILReader() self.assertIsInstance(inst, PILReader) + inst = NrrdReader() + self.assertIsInstance(inst, NrrdReader) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_integration_bundle_run.py b/tests/test_integration_bundle_run.py index 7579858bbc..6f20c55fe2 100644 --- a/tests/test_integration_bundle_run.py +++ b/tests/test_integration_bundle_run.py @@ -48,8 +48,14 @@ def tearDown(self): def test_tiny(self): config_file = os.path.join(self.data_dir, "tiny_config.json") with open(config_file, "w") as f: - json.dump({"": {"_target_": "tests.test_integration_bundle_run._Runnable42", "val": 42}}, f) - cmd = ["coverage", "run", "-m", "monai.bundle", "run", "--config_file", config_file] + json.dump( + { + "trainer": {"_target_": "tests.test_integration_bundle_run._Runnable42", "val": 42}, + "training": "$@trainer.run()", + }, + f, + ) + cmd = ["coverage", "run", "-m", "monai.bundle", "run", "training", "--config_file", config_file] subprocess.check_call(cmd) @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) @@ -81,26 +87,25 @@ def test_shape(self, config_file, expected_shape): # test override with the whole overriding file json.dump("Dataset", f) - saver = LoadImage(image_only=True) - if sys.platform == "win32": override = "--network $@network_def.to(@device) --dataset#_target_ Dataset" else: override = f"--network %{overridefile1}#move_net --dataset#_target_ %{overridefile2}" # test with `monai.bundle` as CLI entry directly - cmd = f"-m monai.bundle run evaluator --postprocessing#transforms#2#output_postfix seg {override}" + cmd = f"-m monai.bundle run evaluating --postprocessing#transforms#2#output_postfix seg {override}" la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file] test_env = os.environ.copy() print(f"CUDA_VISIBLE_DEVICES in {__file__}", test_env.get("CUDA_VISIBLE_DEVICES")) subprocess.check_call(la + ["--args_file", def_args_file], env=test_env) - self.assertTupleEqual(saver(os.path.join(tempdir, "image", "image_seg.nii.gz")).shape, expected_shape) + loader = LoadImage(image_only=True) + self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_seg.nii.gz")).shape, expected_shape) # here test the script with `google fire` tool as CLI - cmd = "-m fire monai.bundle.scripts run --runner_id evaluator" + cmd = "-m fire monai.bundle.scripts run --runner_id evaluating" cmd += f" --evaluator#amp False {override}" la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file] subprocess.check_call(la, env=test_env) - self.assertTupleEqual(saver(os.path.join(tempdir, "image", "image_trans.nii.gz")).shape, expected_shape) + self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_trans.nii.gz")).shape, expected_shape) if __name__ == "__main__": diff --git a/tests/test_integration_classification_2d.py b/tests/test_integration_classification_2d.py index 5a742ce4f9..78f0bf9f36 100644 --- a/tests/test_integration_classification_2d.py +++ b/tests/test_integration_classification_2d.py @@ -33,7 +33,6 @@ RandRotate, RandZoom, ScaleIntensity, - ToTensor, Transpose, ) from monai.utils import set_determinism @@ -62,22 +61,21 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0", # define transforms for image and classification train_transforms = Compose( [ - LoadImage(image_only=True), + LoadImage(image_only=True, simple_keys=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity(), RandRotate(range_x=np.pi / 12, prob=0.5, keep_size=True, dtype=np.float64), RandFlip(spatial_axis=0, prob=0.5), RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5), - ToTensor(), ] ) train_transforms.set_random_state(1234) val_transforms = Compose( - [LoadImage(image_only=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity(), ToTensor()] + [LoadImage(image_only=True, simple_keys=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity()] ) - y_pred_trans = Compose([ToTensor(), Activations(softmax=True)]) - y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=len(np.unique(train_y)))]) + y_pred_trans = Compose([Activations(softmax=True)]) + y_trans = AsDiscrete(to_onehot=len(np.unique(train_y))) auc_metric = ROCAUCMetric() # create train, val data loaders @@ -132,7 +130,7 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0", acc_metric = acc_value.sum().item() / len(acc_value) # decollate prediction and label and execute post processing y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)] - y = [y_trans(i) for i in decollate_batch(y)] + y = [y_trans(i) for i in decollate_batch(y, detach=False)] # compute AUC auc_metric(y_pred, y) auc_value = auc_metric.aggregate() @@ -153,7 +151,7 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0", def run_inference_test(root_dir, test_x, test_y, device="cuda:0", num_workers=10): # define transforms for image and classification - val_transforms = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity(), ToTensor()]) + val_transforms = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity()]) val_ds = MedNISTDataset(test_x, test_y, val_transforms) val_loader = DataLoader(val_ds, batch_size=300, num_workers=num_workers) diff --git a/tests/test_integration_determinism.py b/tests/test_integration_determinism.py index 64c018b4f5..94d2325514 100644 --- a/tests/test_integration_determinism.py +++ b/tests/test_integration_determinism.py @@ -18,7 +18,7 @@ from monai.data import create_test_image_2d from monai.losses import DiceLoss from monai.networks.nets import UNet -from monai.transforms import AddChannel, Compose, RandRotate90, RandSpatialCrop, ScaleIntensity, ToTensor +from monai.transforms import AddChannel, Compose, RandRotate90, RandSpatialCrop, ScaleIntensity from monai.utils import set_determinism from tests.utils import DistTestCase, TimedCall @@ -47,7 +47,7 @@ def __len__(self): loss = DiceLoss(sigmoid=True) opt = torch.optim.Adam(net.parameters(), 1e-2) train_transforms = Compose( - [AddChannel(), ScaleIntensity(), RandSpatialCrop((96, 96), random_size=False), RandRotate90(), ToTensor()] + [AddChannel(), ScaleIntensity(), RandSpatialCrop((96, 96), random_size=False), RandRotate90()] ) src = DataLoader(_TestBatch(train_transforms), batch_size=batch_size, shuffle=True) diff --git a/tests/test_integration_fast_train.py b/tests/test_integration_fast_train.py index 51b2ac1d3f..13f918d201 100644 --- a/tests/test_integration_fast_train.py +++ b/tests/test_integration_fast_train.py @@ -34,8 +34,6 @@ Compose, CropForegroundd, EnsureChannelFirstd, - EnsureType, - EnsureTyped, FgBgToIndicesd, LoadImaged, RandAffined, @@ -52,12 +50,11 @@ ToDeviced, ) from monai.utils import set_determinism -from tests.utils import DistTestCase, SkipIfBeforePyTorchVersion, TimedCall, skip_if_no_cuda, skip_if_quick +from tests.utils import DistTestCase, TimedCall, skip_if_no_cuda, skip_if_quick @skip_if_no_cuda @skip_if_quick -@SkipIfBeforePyTorchVersion((1, 7)) class IntegrationFastTrain(DistTestCase): def setUp(self): set_determinism(seed=0) @@ -95,8 +92,6 @@ def test_train_timing(self): # pre-compute foreground and background indexes # and cache them to accelerate training FgBgToIndicesd(keys="label", fg_postfix="_fg", bg_postfix="_bg"), - # change to execute transforms with Tensor data - EnsureTyped(keys=["image", "label"]), # move the data to GPU and cache to avoid CPU -> GPU sync in every epoch ToDeviced(keys=["image", "label"], device=device), # randomly crop out patch samples from big @@ -138,7 +133,6 @@ def test_train_timing(self): Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")), ScaleIntensityd(keys="image"), CropForegroundd(keys=["image", "label"], source_key="image"), - EnsureTyped(keys=["image", "label"]), # move the data to GPU and cache to avoid CPU -> GPU sync in every epoch ToDeviced(keys=["image", "label"], device=device), ] @@ -171,8 +165,8 @@ def test_train_timing(self): optimizer = Novograd(model.parameters(), learning_rate * 10) scaler = torch.cuda.amp.GradScaler() - post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=2)]) - post_label = Compose([EnsureType(), AsDiscrete(to_onehot=2)]) + post_pred = AsDiscrete(argmax=True, to_onehot=2) + post_label = AsDiscrete(to_onehot=2) dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False) diff --git a/tests/test_integration_segmentation_3d.py b/tests/test_integration_segmentation_3d.py index 5c273d0a46..b5b1d69565 100644 --- a/tests/test_integration_segmentation_3d.py +++ b/tests/test_integration_segmentation_3d.py @@ -18,7 +18,6 @@ import nibabel as nib import numpy as np import torch -from torch.utils.tensorboard import SummaryWriter import monai from monai.data import create_test_image_3d, decollate_batch @@ -37,15 +36,14 @@ SaveImage, ScaleIntensityd, Spacingd, - ToTensor, - ToTensord, ) -from monai.utils import set_determinism -from monai.utils.enums import PostFix +from monai.utils import optional_import, set_determinism from monai.visualize import plot_2d_or_3d_image from tests.testing_data.integration_answers import test_integration_value from tests.utils import DistTestCase, TimedCall, skip_if_quick +SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") + TASK = "integration_segmentation_3d" @@ -69,7 +67,6 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None, keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4 ), RandRotate90d(keys=["img", "seg"], prob=0.8, spatial_axes=[0, 2]), - ToTensord(keys=["img", "seg"]), ] ) train_transforms.set_random_state(1234) @@ -81,7 +78,6 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None, # slight different results between PyTorch 1.5 an 1.6 Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32), ScaleIntensityd(keys="img"), - ToTensord(keys=["img", "seg"]), ] ) @@ -97,7 +93,7 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None, # create a validation data loader val_ds = monai.data.Dataset(data=val_files, transform=val_transforms) val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4) - val_post_tran = Compose([ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) + val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False) # create UNet, DiceLoss and Adam optimizer @@ -182,6 +178,13 @@ def run_inference_test(root_dir, device="cuda:0"): segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz"))) val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)] + saver = SaveImage( + output_dir=os.path.join(root_dir, "output"), + dtype=np.float32, + output_ext=".nii.gz", + output_postfix="seg", + mode="bilinear", + ) # define transforms for image and segmentation val_transforms = Compose( [ @@ -191,13 +194,12 @@ def run_inference_test(root_dir, device="cuda:0"): # slight different results between PyTorch 1.5 an 1.6 Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32), ScaleIntensityd(keys="img"), - ToTensord(keys=["img", "seg"]), ] ) val_ds = monai.data.Dataset(data=val_files, transform=val_transforms) # sliding window inference need to input 1 image in every iteration val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4) - val_post_tran = Compose([ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) + val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5), saver]) dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False) model = UNet( @@ -214,13 +216,6 @@ def run_inference_test(root_dir, device="cuda:0"): with eval_mode(model): # resampling with align_corners=True or dtype=float64 will generate # slight different results between PyTorch 1.5 an 1.6 - saver = SaveImage( - output_dir=os.path.join(root_dir, "output"), - dtype=np.float32, - output_ext=".nii.gz", - output_postfix="seg", - mode="bilinear", - ) for val_data in val_loader: val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device) # define sliding window size and batch size for windows inference @@ -228,11 +223,8 @@ def run_inference_test(root_dir, device="cuda:0"): val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model) # decollate prediction into a list val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)] - val_meta = decollate_batch(val_data[PostFix.meta("img")]) # compute metrics dice_metric(y_pred=val_outputs, y=val_labels) - for img, meta in zip(val_outputs, val_meta): # save a decollated batch of files - saver(img, meta) return dice_metric.aggregate().item() diff --git a/tests/test_integration_sliding_window.py b/tests/test_integration_sliding_window.py index af49e3db77..9b1c7e5200 100644 --- a/tests/test_integration_sliding_window.py +++ b/tests/test_integration_sliding_window.py @@ -16,21 +16,20 @@ import nibabel as nib import numpy as np import torch -from ignite.engine import Engine +from ignite.engine import Engine, Events from torch.utils.data import DataLoader from monai.data import ImageDataset, create_test_image_3d -from monai.handlers import SegmentationSaver from monai.inferers import sliding_window_inference from monai.networks import eval_mode, predict_segmentation from monai.networks.nets import UNet -from monai.transforms import AddChannel -from monai.utils import set_determinism +from monai.transforms import AddChannel, SaveImage +from monai.utils import pytorch_after, set_determinism from tests.utils import DistTestCase, TimedCall, make_nifti_image, skip_if_quick def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"): - ds = ImageDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=False) + ds = ImageDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=True) loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available()) net = UNet( @@ -39,18 +38,23 @@ def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"): roi_size = (16, 32, 48) sw_batch_size = batch_size + saver = SaveImage(output_dir=output_dir, output_ext=".nii.gz", output_postfix="seg") + def _sliding_window_processor(_engine, batch): img = batch[0] # first item from ImageDataset is the input image with eval_mode(net): seg_probs = sliding_window_inference(img.to(device), roi_size, sw_batch_size, net, device=device) return predict_segmentation(seg_probs) - infer_engine = Engine(_sliding_window_processor) - - SegmentationSaver( # 3rd item for image batch meta data - output_dir=output_dir, output_ext=".nii.gz", output_postfix="seg", batch_transform=lambda x: x[2] - ).attach(infer_engine) + def save_func(engine): + if pytorch_after(1, 9, 1): + for m in engine.state.output: + saver(m) + else: + saver(engine.state.output[0]) + infer_engine = Engine(_sliding_window_processor) + infer_engine.add_event_handler(Events.ITERATION_COMPLETED, save_func) infer_engine.run(loader) basename = os.path.basename(img_name)[: -len(".nii.gz")] diff --git a/tests/test_integration_workflows.py b/tests/test_integration_workflows.py index 432e5e90a0..0ef95d4005 100644 --- a/tests/test_integration_workflows.py +++ b/tests/test_integration_workflows.py @@ -19,8 +19,8 @@ import nibabel as nib import numpy as np import torch +from ignite.engine import Events from ignite.metrics import Accuracy -from torch.utils.tensorboard import SummaryWriter import monai from monai.data import create_test_image_3d @@ -30,7 +30,6 @@ CheckpointSaver, LrScheduleHandler, MeanDice, - SegmentationSaver, StatsHandler, TensorBoardImageHandler, TensorBoardStatsHandler, @@ -47,14 +46,15 @@ LoadImaged, RandCropByPosNegLabeld, RandRotate90d, + SaveImage, SaveImaged, ScaleIntensityd, - ToTensord, ) -from monai.utils import set_determinism -from monai.utils.enums import PostFix +from monai.utils import optional_import, set_determinism from tests.testing_data.integration_answers import test_integration_value -from tests.utils import DistTestCase, TimedCall, skip_if_quick +from tests.utils import DistTestCase, TimedCall, pytorch_after, skip_if_quick + +SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter") TASK = "integration_workflows" @@ -75,7 +75,6 @@ def run_training_test(root_dir, device="cuda:0", amp=False, num_workers=4): keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4 ), RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]), - ToTensord(keys=["image", "label"]), ] ) val_transforms = Compose( @@ -83,7 +82,6 @@ def run_training_test(root_dir, device="cuda:0", amp=False, num_workers=4): LoadImaged(keys=["image", "label"]), AsChannelFirstd(keys=["image", "label"], channel_dim=-1), ScaleIntensityd(keys=["image", "label"]), - ToTensord(keys=["image", "label"]), ] ) @@ -111,7 +109,6 @@ def run_training_test(root_dir, device="cuda:0", amp=False, num_workers=4): val_postprocessing = Compose( [ - ToTensord(keys=["pred", "label"]), Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold=0.5), KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), @@ -147,12 +144,13 @@ def _forward_completed(self, engine): additional_metrics={"val_acc": Accuracy(output_transform=from_engine(["pred", "label"]))}, metric_cmp_fn=lambda cur, prev: cur >= prev, # if greater or equal, treat as new best metric val_handlers=val_handlers, - amp=True if amp else False, + amp=bool(amp), + to_kwargs={"memory_format": torch.preserve_format}, + amp_kwargs={"dtype": torch.float16 if bool(amp) else torch.float32} if pytorch_after(1, 10, 0) else {}, ) train_postprocessing = Compose( [ - ToTensord(keys=["pred", "label"]), Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold=0.5), KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), @@ -200,8 +198,10 @@ def _model_completed(self, engine): postprocessing=train_postprocessing, key_train_metric={"train_acc": Accuracy(output_transform=from_engine(["pred", "label"]))}, train_handlers=train_handlers, - amp=True if amp else False, + amp=bool(amp), optim_set_to_none=True, + to_kwargs={"memory_format": torch.preserve_format}, + amp_kwargs={"dtype": torch.float16 if bool(amp) else torch.float32} if pytorch_after(1, 10, 0) else {}, ) trainer.run() @@ -219,7 +219,6 @@ def run_inference_test(root_dir, model_file, device="cuda:0", amp=False, num_wor LoadImaged(keys=["image", "label"]), AsChannelFirstd(keys=["image", "label"], channel_dim=-1), ScaleIntensityd(keys=["image", "label"]), - ToTensord(keys=["image", "label"]), ] ) @@ -239,27 +238,24 @@ def run_inference_test(root_dir, model_file, device="cuda:0", amp=False, num_wor val_postprocessing = Compose( [ - ToTensord(keys=["pred", "label"]), Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold=0.5), KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), # test the case that `pred` in `engine.state.output`, while `image_meta_dict` in `engine.state.batch` - SaveImaged( - keys="pred", meta_keys=PostFix.meta("image"), output_dir=root_dir, output_postfix="seg_transform" - ), + SaveImaged(keys="pred", output_dir=root_dir, output_postfix="seg_transform"), ] ) val_handlers = [ StatsHandler(iteration_log=False), CheckpointLoader(load_path=f"{model_file}", load_dict={"net": net}), - SegmentationSaver( - output_dir=root_dir, - output_postfix="seg_handler", - batch_transform=from_engine(PostFix.meta("image")), - output_transform=from_engine("pred"), - ), ] + saver = SaveImage(output_dir=root_dir, output_postfix="seg_handler") + + def save_func(engine): + for o in from_engine("pred")(engine.state.output): + saver(o) + evaluator = SupervisedEvaluator( device=device, val_data_loader=val_loader, @@ -271,8 +267,9 @@ def run_inference_test(root_dir, model_file, device="cuda:0", amp=False, num_wor }, additional_metrics={"val_acc": Accuracy(output_transform=from_engine(["pred", "label"]))}, val_handlers=val_handlers, - amp=True if amp else False, + amp=bool(amp), ) + evaluator.add_event_handler(Events.ITERATION_COMPLETED, save_func) evaluator.run() return evaluator.state.best_metric diff --git a/tests/test_integration_workflows_gan.py b/tests/test_integration_workflows_gan.py index c9306b349f..ff53851ce0 100644 --- a/tests/test_integration_workflows_gan.py +++ b/tests/test_integration_workflows_gan.py @@ -26,7 +26,7 @@ from monai.handlers import CheckpointSaver, StatsHandler, TensorBoardStatsHandler from monai.networks import normal_init from monai.networks.nets import Discriminator, Generator -from monai.transforms import AsChannelFirstd, Compose, LoadImaged, RandFlipd, ScaleIntensityd, ToTensord +from monai.transforms import AsChannelFirstd, Compose, LoadImaged, RandFlipd, ScaleIntensityd from monai.utils import set_determinism from tests.utils import DistTestCase, TimedCall, skip_if_quick @@ -42,7 +42,6 @@ def run_training_test(root_dir, device="cuda:0"): AsChannelFirstd(keys=["reals"]), ScaleIntensityd(keys=["reals"]), RandFlipd(keys=["reals"], prob=0.5), - ToTensord(keys=["reals"]), ] ) train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5) @@ -117,6 +116,7 @@ def generator_loss(gen_images): latent_shape=latent_size, key_train_metric=key_train_metric, train_handlers=train_handlers, + to_kwargs={"memory_format": torch.preserve_format, "dtype": torch.float32}, ) trainer.run() diff --git a/tests/test_inverse.py b/tests/test_inverse.py index c04e9b0cd7..82902a09eb 100644 --- a/tests/test_inverse.py +++ b/tests/test_inverse.py @@ -20,13 +20,11 @@ import torch from parameterized import parameterized -from monai.data import CacheDataset, DataLoader, create_test_image_2d, create_test_image_3d -from monai.data.utils import decollate_batch +from monai.data import CacheDataset, DataLoader, MetaTensor, create_test_image_2d, create_test_image_3d, decollate_batch from monai.networks.nets import UNet from monai.transforms import ( AddChanneld, Affined, - BatchInverseTransform, BorderPadd, CenterScaleCropd, CenterSpatialCropd, @@ -34,6 +32,7 @@ CropForegroundd, DivisiblePadd, Flipd, + FromMetaTensord, InvertibleTransform, Lambdad, LoadImaged, @@ -59,11 +58,11 @@ Spacingd, SpatialCropd, SpatialPadd, - TraceableTransform, + ToMetaTensord, Transposed, Zoomd, allow_missing_keys_mode, - convert_inverse_interp_mode, + convert_applied_interp_mode, ) from monai.utils import first, get_seed, optional_import, set_determinism from tests.utils import make_nifti_image, make_rand_affine @@ -97,7 +96,7 @@ partial(RandSpatialCropd, roi_size=12 + val), partial(ResizeWithPadOrCropd, spatial_size=21 - val), ): - TESTS.append((t.func.__name__ + name, name, 0, t(KEYS))) # type: ignore + TESTS.append((t.func.__name__ + name, name, 0, True, t(KEYS))) # type: ignore # non-sensical tests: crop bigger or pad smaller or -ve values for t in ( @@ -110,112 +109,118 @@ partial(SpatialCropd, roi_center=10, roi_size=100), partial(SpatialCropd, roi_start=3, roi_end=100), ): - TESTS.append((t.func.__name__ + "bad 1D even", "1D even", 0, t(KEYS))) # type: ignore + TESTS.append((t.func.__name__ + "bad 1D even", "1D even", 0, True, t(KEYS))) # type: ignore TESTS.append( ( "SpatialPadd (x2) 2d", "2D", 0, + True, SpatialPadd(KEYS, spatial_size=[111, 113], method="end"), SpatialPadd(KEYS, spatial_size=[118, 117]), ) ) -TESTS.append(("SpatialPadd 3d", "3D", 0, SpatialPadd(KEYS, spatial_size=[112, 113, 116]))) +TESTS.append(("SpatialPadd 3d", "3D", 0, True, SpatialPadd(KEYS, spatial_size=[112, 113, 116]))) -TESTS.append(("SpatialCropd 2d", "2D", 0, SpatialCropd(KEYS, [49, 51], [90, 89]))) +TESTS.append(("SpatialCropd 2d", "2D", 0, True, SpatialCropd(KEYS, [49, 51], [90, 89]))) TESTS.append( ( "SpatialCropd 3d", "3D", 0, + True, SpatialCropd(KEYS, roi_slices=[slice(s, e) for s, e in zip([None, None, -99], [None, -2, None])]), ) ) -TESTS.append(("SpatialCropd 2d", "2D", 0, SpatialCropd(KEYS, [49, 51], [390, 89]))) +TESTS.append(("SpatialCropd 2d", "2D", 0, True, SpatialCropd(KEYS, [49, 51], [390, 89]))) -TESTS.append(("SpatialCropd 3d", "3D", 0, SpatialCropd(KEYS, [49, 51, 44], [90, 89, 93]))) +TESTS.append(("SpatialCropd 3d", "3D", 0, True, SpatialCropd(KEYS, [49, 51, 44], [90, 89, 93]))) -TESTS.append(("RandSpatialCropd 2d", "2D", 0, RandSpatialCropd(KEYS, [96, 93], None, True, False))) +TESTS.append(("RandSpatialCropd 2d", "2D", 0, True, RandSpatialCropd(KEYS, [96, 93], None, True, False))) -TESTS.append(("RandSpatialCropd 3d", "3D", 0, RandSpatialCropd(KEYS, [96, 93, 92], None, False, False))) +TESTS.append(("RandSpatialCropd 3d", "3D", 0, True, RandSpatialCropd(KEYS, [96, 93, 92], None, False, False))) -TESTS.append(("BorderPadd 2d", "2D", 0, BorderPadd(KEYS, [3, 7, 2, 5]))) +TESTS.append(("BorderPadd 2d", "2D", 0, True, BorderPadd(KEYS, [3, 7, 2, 5]))) -TESTS.append(("BorderPadd 2d", "2D", 0, BorderPadd(KEYS, [3, 7]))) +TESTS.append(("BorderPadd 2d", "2D", 0, True, BorderPadd(KEYS, [3, 7]))) -TESTS.append(("BorderPadd 3d", "3D", 0, BorderPadd(KEYS, [4]))) +TESTS.append(("BorderPadd 3d", "3D", 0, True, BorderPadd(KEYS, [4]))) -TESTS.append(("DivisiblePadd 2d", "2D", 0, DivisiblePadd(KEYS, k=4))) +TESTS.append(("DivisiblePadd 2d", "2D", 0, True, DivisiblePadd(KEYS, k=4))) -TESTS.append(("DivisiblePadd 3d", "3D", 0, DivisiblePadd(KEYS, k=[4, 8, 11]))) +TESTS.append(("DivisiblePadd 3d", "3D", 0, True, DivisiblePadd(KEYS, k=[4, 8, 11]))) +TESTS.append(("CenterSpatialCropd 2d", "2D", 0, True, CenterSpatialCropd(KEYS, roi_size=95))) -TESTS.append(("CenterSpatialCropd 2d", "2D", 0, CenterSpatialCropd(KEYS, roi_size=95))) +TESTS.append(("CenterSpatialCropd 3d", "3D", 0, True, CenterSpatialCropd(KEYS, roi_size=[95, 97, 98]))) -TESTS.append(("CenterSpatialCropd 3d", "3D", 0, CenterSpatialCropd(KEYS, roi_size=[95, 97, 98]))) +TESTS.append(("CropForegroundd 2d", "2D", 0, True, CropForegroundd(KEYS, source_key="label", margin=2))) -TESTS.append(("CropForegroundd 2d", "2D", 0, CropForegroundd(KEYS, source_key="label", margin=2))) +TESTS.append(("CropForegroundd 3d", "3D", 0, True, CropForegroundd(KEYS, source_key="label", k_divisible=[5, 101, 2]))) -TESTS.append(("CropForegroundd 3d", "3D", 0, CropForegroundd(KEYS, source_key="label", k_divisible=[5, 101, 2]))) +TESTS.append(("ResizeWithPadOrCropd 3d", "3D", 0, True, ResizeWithPadOrCropd(KEYS, [201, 150, 105]))) +TESTS.append(("Flipd 3d", "3D", 0, True, Flipd(KEYS, [1, 2]))) +TESTS.append(("Flipd 3d", "3D", 0, True, Flipd(KEYS, [1, 2]))) -TESTS.append(("ResizeWithPadOrCropd 3d", "3D", 0, ResizeWithPadOrCropd(KEYS, [201, 150, 105]))) +TESTS.append(("RandFlipd 3d", "3D", 0, True, RandFlipd(KEYS, 1, [1, 2]))) -TESTS.append(("Flipd 3d", "3D", 0, Flipd(KEYS, [1, 2]))) - -TESTS.append(("RandFlipd 3d", "3D", 0, RandFlipd(KEYS, 1, [1, 2]))) - -TESTS.append(("RandAxisFlipd 3d", "3D", 0, RandAxisFlipd(KEYS, 1))) +TESTS.append(("RandAxisFlipd 3d", "3D", 0, True, RandAxisFlipd(KEYS, 1))) +TESTS.append(("RandAxisFlipd 3d", "3D", 0, True, RandAxisFlipd(KEYS, 1))) for acc in [True, False]: - TESTS.append(("Orientationd 3d", "3D", 0, Orientationd(KEYS, "RAS", as_closest_canonical=acc))) + TESTS.append(("Orientationd 3d", "3D", 0, True, Orientationd(KEYS, "RAS", as_closest_canonical=acc))) -TESTS.append(("Rotate90d 2d", "2D", 0, Rotate90d(KEYS))) +TESTS.append(("Rotate90d 2d", "2D", 0, True, Rotate90d(KEYS))) -TESTS.append(("Rotate90d 3d", "3D", 0, Rotate90d(KEYS, k=2, spatial_axes=(1, 2)))) +TESTS.append(("Rotate90d 3d", "3D", 0, True, Rotate90d(KEYS, k=2, spatial_axes=(1, 2)))) -TESTS.append(("RandRotate90d 3d", "3D", 0, RandRotate90d(KEYS, prob=1, spatial_axes=(1, 2)))) +TESTS.append(("RandRotate90d 3d", "3D", 0, True, RandRotate90d(KEYS, prob=1, spatial_axes=(1, 2)))) -TESTS.append(("Spacingd 3d", "3D", 3e-2, Spacingd(KEYS, [0.5, 0.7, 0.9], diagonal=False))) +TESTS.append(("Spacingd 3d", "3D", 3e-2, True, Spacingd(KEYS, [0.5, 0.7, 0.9], diagonal=False))) -TESTS.append(("Resized 2d", "2D", 2e-1, Resized(KEYS, [50, 47]))) +TESTS.append(("Resized 2d", "2D", 2e-1, True, Resized(KEYS, [50, 47]))) -TESTS.append(("Resized 3d", "3D", 5e-2, Resized(KEYS, [201, 150, 78]))) +TESTS.append(("Resized 3d", "3D", 5e-2, True, Resized(KEYS, [201, 150, 78]))) -TESTS.append(("Resized longest 2d", "2D", 2e-1, Resized(KEYS, 47, "longest", "area"))) +TESTS.append(("Resized longest 2d", "2D", 2e-1, True, Resized(KEYS, 47, "longest", "area"))) -TESTS.append(("Resized longest 3d", "3D", 5e-2, Resized(KEYS, 201, "longest", "trilinear", True))) +TESTS.append(("Resized longest 3d", "3D", 5e-2, True, Resized(KEYS, 201, "longest", "trilinear", True))) -TESTS.append(("Lambdad 2d", "2D", 5e-2, Lambdad(KEYS, func=lambda x: x + 5, inv_func=lambda x: x - 5, overwrite=True))) +TESTS.append( + ("Lambdad 2d", "2D", 5e-2, False, Lambdad(KEYS, func=lambda x: x + 5, inv_func=lambda x: x - 5, overwrite=True)) +) TESTS.append( ( "RandLambdad 3d", "3D", 5e-2, + False, RandLambdad(KEYS, func=lambda x: x * 10, inv_func=lambda x: x / 10, overwrite=True, prob=0.5), ) ) -TESTS.append(("Zoomd 1d", "1D odd", 0, Zoomd(KEYS, zoom=2, keep_size=False))) +TESTS.append(("Zoomd 1d", "1D odd", 0, True, Zoomd(KEYS, zoom=2, keep_size=False))) -TESTS.append(("Zoomd 2d", "2D", 2e-1, Zoomd(KEYS, zoom=0.9))) +TESTS.append(("Zoomd 2d", "2D", 2e-1, True, Zoomd(KEYS, zoom=0.9))) -TESTS.append(("Zoomd 3d", "3D", 3e-2, Zoomd(KEYS, zoom=[2.5, 1, 3], keep_size=False))) +TESTS.append(("Zoomd 3d", "3D", 3e-2, True, Zoomd(KEYS, zoom=[2.5, 1, 3], keep_size=False))) -TESTS.append(("RandZoom 3d", "3D", 9e-2, RandZoomd(KEYS, 1, [0.5, 0.6, 0.9], [1.1, 1, 1.05], keep_size=True))) +TESTS.append(("RandZoom 3d", "3D", 9e-2, True, RandZoomd(KEYS, 1, [0.5, 0.6, 0.9], [1.1, 1, 1.05], keep_size=True))) -TESTS.append(("RandRotated, prob 0", "2D", 0, RandRotated(KEYS, prob=0, dtype=np.float64))) +TESTS.append(("RandRotated, prob 0", "2D", 0, True, RandRotated(KEYS, prob=0, dtype=np.float64))) TESTS.append( ( "Rotated 2d", "2D", 8e-2, + True, Rotated(KEYS, random.uniform(np.pi / 6, np.pi), keep_size=True, align_corners=False, dtype=np.float64), ) ) @@ -225,6 +230,7 @@ "Rotated 3d", "3D", 1e-1, + True, Rotated(KEYS, [random.uniform(np.pi / 6, np.pi) for _ in range(3)], True, dtype=np.float64), ) ) @@ -234,19 +240,21 @@ "RandRotated 3d", "3D", 1e-1, + True, RandRotated(KEYS, *[random.uniform(np.pi / 6, np.pi) for _ in range(3)], 1, dtype=np.float64), # type: ignore ) ) -TESTS.append(("Transposed 2d", "2D", 0, Transposed(KEYS, [0, 2, 1]))) # channel=0 +TESTS.append(("Transposed 2d", "2D", 0, False, Transposed(KEYS, [0, 2, 1]))) # channel=0 -TESTS.append(("Transposed 3d", "3D", 0, Transposed(KEYS, [0, 3, 1, 2]))) # channel=0 +TESTS.append(("Transposed 3d", "3D", 0, False, Transposed(KEYS, [0, 3, 1, 2]))) # channel=0 TESTS.append( ( "Affine 3d", "3D", 1e-1, + True, Affined( KEYS, spatial_size=[155, 179, 192], @@ -263,6 +271,7 @@ "RandAffine 3d", "3D", 1e-1, + True, RandAffined( KEYS, [155, 179, 192], @@ -276,24 +285,27 @@ ) ) -TESTS.append(("RandAffine 3d", "3D", 0, RandAffined(KEYS, spatial_size=None, prob=0))) +TESTS.append(("RandAffine 3d", "3D", 0, True, RandAffined(KEYS, spatial_size=None, prob=0))) TESTS.append( ( "RandCropByLabelClassesd 2d", "2D", 1e-7, + True, RandCropByLabelClassesd(KEYS, "label", (99, 96), ratios=[1, 2, 3, 4, 5], num_classes=5, num_samples=10), ) ) -TESTS.append(("RandCropByPosNegLabeld 2d", "2D", 1e-7, RandCropByPosNegLabeld(KEYS, "label", (99, 96), num_samples=10))) +TESTS.append( + ("RandCropByPosNegLabeld 2d", "2D", 1e-7, True, RandCropByPosNegLabeld(KEYS, "label", (99, 96), num_samples=10)) +) -TESTS.append(("RandSpatialCropSamplesd 2d", "2D", 1e-7, RandSpatialCropSamplesd(KEYS, (90, 91), num_samples=10))) +TESTS.append(("RandSpatialCropSamplesd 2d", "2D", 1e-7, True, RandSpatialCropSamplesd(KEYS, (90, 91), num_samples=10))) -TESTS.append(("RandWeightedCropd 2d", "2D", 1e-7, RandWeightedCropd(KEYS, "label", (90, 91), num_samples=10))) +TESTS.append(("RandWeightedCropd 2d", "2D", 1e-7, True, RandWeightedCropd(KEYS, "label", (90, 91), num_samples=10))) -TESTS_COMPOSE_X2 = [(t[0] + " Compose", t[1], t[2], Compose(Compose(t[3:]))) for t in TESTS] +TESTS_COMPOSE_X2 = [(t[0] + " Compose", t[1], t[2], t[3], Compose(Compose(t[4:]))) for t in TESTS] TESTS = TESTS + TESTS_COMPOSE_X2 # type: ignore @@ -365,15 +377,15 @@ def setUp(self): im_1d = np.pad(np.arange(size), 5)[None] name = "1D even" if size % 2 == 0 else "1D odd" self.all_data[name] = { - "image": np.array(im_1d, copy=True), - "label": np.array(im_1d, copy=True), - "other": np.array(im_1d, copy=True), + "image": torch.as_tensor(np.array(im_1d, copy=True)), + "label": torch.as_tensor(np.array(im_1d, copy=True)), + "other": torch.as_tensor(np.array(im_1d, copy=True)), } im_2d_fname, seg_2d_fname = (make_nifti_image(i) for i in create_test_image_2d(101, 100)) im_3d_fname, seg_3d_fname = (make_nifti_image(i, affine) for i in create_test_image_3d(100, 101, 107)) - load_ims = Compose([LoadImaged(KEYS), AddChanneld(KEYS)]) + load_ims = Compose([LoadImaged(KEYS), AddChanneld(KEYS), FromMetaTensord(KEYS)]) self.all_data["2D"] = load_ims({"image": im_2d_fname, "label": seg_2d_fname}) self.all_data["3D"] = load_ims({"image": im_3d_fname, "label": seg_3d_fname}) @@ -406,10 +418,12 @@ def check_inverse(self, name, keys, orig_d, fwd_bck_d, unmodified_d, acceptable_ raise @parameterized.expand(TESTS) - def test_inverse(self, _, data_name, acceptable_diff, *transforms): + def test_inverse(self, _, data_name, acceptable_diff, is_meta, *transforms): name = _ data = self.all_data[data_name] + if is_meta: + data = ToMetaTensord(KEYS)(data) forwards = [data.copy()] @@ -457,40 +471,42 @@ def test_inverse_inferred_seg(self, extra_transform): # num workers = 0 for mac num_workers = 2 if sys.platform == "linux" else 0 transforms = Compose([AddChanneld(KEYS), SpatialPadd(KEYS, (150, 153)), extra_transform]) - num_invertible_transforms = sum(1 for i in transforms.transforms if isinstance(i, InvertibleTransform)) dataset = CacheDataset(test_data, transform=transforms, progress=False) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) device = "cuda" if torch.cuda.is_available() else "cpu" - model = UNet(spatial_dims=2, in_channels=1, out_channels=1, channels=(2, 4), strides=(2,)).to(device) + model = UNet(spatial_dims=2, in_channels=1, out_channels=1, channels=(2, 4), strides=(1,)).to(device) data = first(loader) - self.assertEqual(len(data["label_transforms"]), num_invertible_transforms) self.assertEqual(data["image"].shape[0], batch_size * NUM_SAMPLES) labels = data["label"].to(device) + self.assertIsInstance(labels, MetaTensor) segs = model(labels).detach().cpu() - label_transform_key = TraceableTransform.trace_key("label") - segs_dict = {"label": segs, label_transform_key: data[label_transform_key]} - - segs_dict_decollated = decollate_batch(segs_dict) + segs_decollated = decollate_batch(segs) + self.assertIsInstance(segs_decollated[0], MetaTensor) # inverse of individual segmentation - seg_dict = first(segs_dict_decollated) + seg_metatensor = first(segs_decollated) # test to convert interpolation mode for 1 data of model output batch - convert_inverse_interp_mode(seg_dict, mode="nearest", align_corners=None) + convert_applied_interp_mode(seg_metatensor.applied_operations, mode="nearest", align_corners=None) + + # manually invert the last crop samples + xform = seg_metatensor.applied_operations.pop(-1) + shape_before_extra_xform = xform["orig_size"] + resizer = ResizeWithPadOrCrop(spatial_size=shape_before_extra_xform) + with resizer.trace_transform(False): + seg_metatensor = resizer(seg_metatensor) with allow_missing_keys_mode(transforms): - inv_seg = transforms.inverse(seg_dict)["label"] - self.assertEqual(len(data["label_transforms"]), num_invertible_transforms) - self.assertEqual(len(seg_dict["label_transforms"]), num_invertible_transforms) + inv_seg = transforms.inverse({"label": seg_metatensor})["label"] self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape) - # Inverse of batch - batch_inverter = BatchInverseTransform(transforms, loader, collate_fn=no_collation, detach=True) - with allow_missing_keys_mode(transforms): - inv_batch = batch_inverter(segs_dict) - self.assertEqual(inv_batch[0]["label"].shape[1:], test_data[0]["label"].shape) + # # Inverse of batch + # batch_inverter = BatchInverseTransform(transforms, loader, collate_fn=no_collation, detach=True) + # with allow_missing_keys_mode(transforms): + # inv_batch = batch_inverter(first(loader)) + # self.assertEqual(inv_batch[0]["label"].shape[1:], test_data[0]["label"].shape) if __name__ == "__main__": diff --git a/tests/test_inverse_array.py b/tests/test_inverse_array.py new file mode 100644 index 0000000000..bb9669e8df --- /dev/null +++ b/tests/test_inverse_array.py @@ -0,0 +1,67 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.data import MetaTensor +from monai.transforms import AddChannel, Compose, Flip, Orientation, Spacing +from monai.transforms.inverse import InvertibleTransform +from monai.utils import optional_import +from tests.utils import TEST_DEVICES + +_, has_nib = optional_import("nibabel") + +TESTS = [] +for use_compose in (False, True): + for dtype in (torch.float32, torch.float64): + for device in TEST_DEVICES: + TESTS.append([use_compose, dtype, *device]) + + +@unittest.skipUnless(has_nib, "Requires nibabel") +class TestInverseArray(unittest.TestCase): + @staticmethod + def get_image(dtype, device) -> MetaTensor: + affine = torch.tensor([[0, 0, 1, 0], [-1, 0, 0, 0], [0, 10, 0, 0], [0, 0, 0, 1]]).to(dtype).to(device) + img = torch.rand((15, 16, 17)).to(dtype).to(device) + return MetaTensor(img, affine=affine) + + @parameterized.expand(TESTS) + def test_inverse_array(self, use_compose, dtype, device): + img: MetaTensor + tr = Compose([AddChannel(), Orientation("RAS"), Flip(1), Spacing([1.0, 1.2, 0.9], align_corners=False)]) + num_invertible = len([i for i in tr.transforms if isinstance(i, InvertibleTransform)]) + + # forward + img = tr(self.get_image(dtype, device)) + self.assertEqual(len(img.applied_operations), num_invertible) + + # inverse with Compose + if use_compose: + img = tr.inverse(img) + self.assertEqual(len(img.applied_operations), 0) + + # inverse individually + else: + _tr: InvertibleTransform + num_to_inverse = num_invertible + for _tr in tr.transforms[::-1]: + if isinstance(_tr, InvertibleTransform): + img = _tr.inverse(img) + num_to_inverse -= 1 + self.assertEqual(len(img.applied_operations), num_to_inverse) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_inverse_collation.py b/tests/test_inverse_collation.py index 4e8c6b58cc..4614432808 100644 --- a/tests/test_inverse_collation.py +++ b/tests/test_inverse_collation.py @@ -17,10 +17,19 @@ import torch from parameterized import parameterized -from monai.data import CacheDataset, DataLoader, create_test_image_2d, create_test_image_3d, pad_list_data_collate +from monai.data import ( + CacheDataset, + DataLoader, + MetaTensor, + create_test_image_2d, + create_test_image_3d, + decollate_batch, + pad_list_data_collate, +) from monai.transforms import ( AddChanneld, Compose, + Flipd, LoadImaged, RandAffined, RandAxisFlipd, @@ -29,7 +38,7 @@ RandRotated, RandZoomd, ResizeWithPadOrCropd, - ToTensord, + Rotated, ) from monai.utils import optional_import, set_determinism from tests.utils import make_nifti_image @@ -46,10 +55,12 @@ (t.__class__.__name__ + (" pad_list_data_collate" if collate_fn else " default_collate"), t, collate_fn, 3) for collate_fn in [None, pad_list_data_collate] for t in [ + Flipd(KEYS, spatial_axis=1), RandFlipd(keys=KEYS, prob=0.5, spatial_axis=[1, 2]), RandAxisFlipd(keys=KEYS, prob=0.5), - Compose([RandRotate90d(keys=KEYS, spatial_axes=(1, 2)), ToTensord(keys=KEYS)]), + Compose([RandRotate90d(keys=KEYS, spatial_axes=(1, 2))]), RandZoomd(keys=KEYS, prob=0.5, min_zoom=0.5, max_zoom=1.1, keep_size=True), + Rotated(keys=KEYS, angle=np.pi, dtype=np.float64), RandRotated(keys=KEYS, prob=0.5, range_x=np.pi, dtype=np.float64), RandAffined( keys=KEYS, prob=0.5, rotate_range=np.pi, device=torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -61,10 +72,12 @@ (t.__class__.__name__ + (" pad_list_data_collate" if collate_fn else " default_collate"), t, collate_fn, 2) for collate_fn in [None, pad_list_data_collate] for t in [ + Flipd(KEYS, spatial_axis=1), RandFlipd(keys=KEYS, prob=0.5, spatial_axis=[1]), RandAxisFlipd(keys=KEYS, prob=0.5), - Compose([RandRotate90d(keys=KEYS, prob=0.5, spatial_axes=(0, 1)), ToTensord(keys=KEYS)]), + Compose([RandRotate90d(keys=KEYS, prob=0.5, spatial_axes=(0, 1))]), RandZoomd(keys=KEYS, prob=0.5, min_zoom=0.5, max_zoom=1.1, keep_size=True), + Rotated(keys=KEYS, angle=np.pi / 2, dtype=np.float64), RandRotated(keys=KEYS, prob=0.5, range_x=np.pi, dtype=np.float64), RandAffined( keys=KEYS, prob=0.5, rotate_range=np.pi, device=torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -99,11 +112,12 @@ def tearDown(self): @parameterized.expand(TESTS_2D + TESTS_3D) def test_collation(self, _, transform, collate_fn, ndim): + """transform, collate_fn, ndim""" data = self.data_3d if ndim == 3 else self.data_2d if collate_fn: modified_transform = transform else: - modified_transform = Compose([transform, ResizeWithPadOrCropd(KEYS, 100), ToTensord(KEYS)]) + modified_transform = Compose([transform, ResizeWithPadOrCropd(KEYS, 100)]) # num workers = 0 for mac or gpu transforms num_workers = 0 if sys.platform != "linux" or torch.cuda.is_available() else 2 @@ -112,9 +126,20 @@ def test_collation(self, _, transform, collate_fn, ndim): loader = DataLoader(dataset, num_workers, batch_size=self.batch_size, collate_fn=collate_fn) for item in loader: - np.testing.assert_array_equal( - item["image_transforms"][0]["do_transforms"], item["label_transforms"][0]["do_transforms"] - ) + if isinstance(item, dict): + np.testing.assert_array_equal(item["image"].shape, item["label"].shape) + continue + d = decollate_batch(item) + self.assertTrue(len(d) <= self.batch_size) + for b in d: + self.assertTrue(isinstance(b["image"], MetaTensor)) + np.testing.assert_array_equal( + b["image"].applied_operations[-1]["orig_size"], b["label"].applied_operations[-1]["orig_size"] + ) + np.testing.assert_array_equal( + b["image"].applied_operations[-1].get("_do_transform"), + b["label"].applied_operations[-1].get("_do_transform"), + ) if __name__ == "__main__": diff --git a/tests/test_invertd.py b/tests/test_invertd.py index 64c26c4012..fc4725d98b 100644 --- a/tests/test_invertd.py +++ b/tests/test_invertd.py @@ -15,13 +15,12 @@ import numpy as np import torch -from monai.data import CacheDataset, DataLoader, create_test_image_3d, decollate_batch +from monai.data import DataLoader, Dataset, create_test_image_3d, decollate_batch from monai.transforms import ( - AddChanneld, CastToTyped, Compose, CopyItemsd, - EnsureTyped, + EnsureChannelFirstd, Invertd, LoadImaged, Orientationd, @@ -34,10 +33,8 @@ ResizeWithPadOrCropd, ScaleIntensityd, Spacingd, - ToTensord, ) from monai.utils import set_determinism -from monai.utils.enums import PostFix from tests.utils import make_nifti_image KEYS = ["image", "label"] @@ -49,23 +46,18 @@ def test_invert(self): im_fname, seg_fname = (make_nifti_image(i) for i in create_test_image_3d(101, 100, 107, noise_max=100)) transform = Compose( [ - LoadImaged(KEYS), - AddChanneld(KEYS), + LoadImaged(KEYS, image_only=True), + EnsureChannelFirstd(KEYS), Orientationd(KEYS, "RPS"), Spacingd(KEYS, pixdim=(1.2, 1.01, 0.9), mode=["bilinear", "nearest"], dtype=np.float32), ScaleIntensityd("image", minv=1, maxv=10), RandFlipd(KEYS, prob=0.5, spatial_axis=[1, 2]), RandAxisFlipd(KEYS, prob=0.5), - RandRotate90d(KEYS, spatial_axes=(1, 2)), + RandRotate90d(KEYS, prob=0, spatial_axes=(1, 2)), RandZoomd(KEYS, prob=0.5, min_zoom=0.5, max_zoom=1.1, keep_size=True), RandRotated(KEYS, prob=0.5, range_x=np.pi, mode="bilinear", align_corners=True, dtype=np.float64), RandAffined(KEYS, prob=0.5, rotate_range=np.pi, mode="nearest"), ResizeWithPadOrCropd(KEYS, 100), - # test EnsureTensor for complicated dict data and invert it - CopyItemsd(PostFix.meta("image"), times=1, names="test_dict"), - # test to support Tensor, Numpy array and dictionary when inverting - EnsureTyped(keys=["image", "test_dict"]), - ToTensord("image"), CastToTyped(KEYS, dtype=[torch.uint8, np.uint8]), CopyItemsd("label", times=2, names=["label_inverted", "label_inverted1"]), CopyItemsd("image", times=2, names=["image_inverted", "image_inverted1"]), @@ -76,17 +68,15 @@ def test_invert(self): # num workers = 0 for mac or gpu transforms num_workers = 0 if sys.platform != "linux" or torch.cuda.is_available() else 2 - dataset = CacheDataset(data, transform=transform, progress=False) - loader = DataLoader(dataset, num_workers=num_workers, batch_size=5) + dataset = Dataset(data, transform=transform) + transform.inverse(dataset[0]) + loader = DataLoader(dataset, num_workers=num_workers, batch_size=1) inverter = Invertd( # `image` was not copied, invert the original value directly - keys=["image_inverted", "label_inverted", "test_dict"], + keys=["image_inverted", "label_inverted"], transform=transform, - orig_keys=["label", "label", "test_dict"], - meta_keys=[PostFix.meta("image_inverted"), PostFix.meta("label_inverted"), None], - orig_meta_keys=[PostFix.meta("label"), PostFix.meta("label"), None], + orig_keys=["label", "label"], nearest_interp=True, - to_tensor=[True, False, False], device="cpu", ) @@ -95,31 +85,11 @@ def test_invert(self): keys=["image_inverted1", "label_inverted1"], transform=transform, orig_keys=["image", "image"], - meta_keys=[PostFix.meta("image_inverted1"), PostFix.meta("label_inverted1")], - orig_meta_keys=[PostFix.meta("image"), PostFix.meta("image")], nearest_interp=[True, False], - to_tensor=[True, True], device="cpu", ) - expected_keys = [ - "image", - "image_inverted", - "image_inverted1", - PostFix.meta("image_inverted1"), - PostFix.meta("image_inverted"), - PostFix.meta("image"), - "image_transforms", - "label", - "label_inverted", - "label_inverted1", - PostFix.meta("label_inverted1"), - PostFix.meta("label_inverted"), - PostFix.meta("label"), - "label_transforms", - "test_dict", - "test_dict_transforms", - ] + expected_keys = ["image", "image_inverted", "image_inverted1", "label", "label_inverted", "label_inverted1"] # execute 1 epoch for d in loader: d = decollate_batch(d) @@ -137,9 +107,6 @@ def test_invert(self): i = item["label_inverted"] torch.testing.assert_allclose(i.to(torch.uint8).to(torch.float), i.to(torch.float)) self.assertTupleEqual(i.shape[1:], (100, 101, 107)) - # test inverted test_dict - self.assertTrue(isinstance(item["test_dict"]["affine"], np.ndarray)) - self.assertTrue(isinstance(item["test_dict"]["filename_or_obj"], str)) # check the case that different items use different interpolation mode to invert transforms d = item["image_inverted1"] @@ -154,16 +121,16 @@ def test_invert(self): # check labels match reverted = item["label_inverted"].detach().cpu().numpy().astype(np.int32) - original = LoadImaged(KEYS)(data[-1])["label"] + original = LoadImaged(KEYS, image_only=True)(data[-1])["label"] n_good = np.sum(np.isclose(reverted, original, atol=1e-3)) - reverted_name = item[PostFix.meta("label_inverted")]["filename_or_obj"] + reverted_name = item["label_inverted"].meta["filename_or_obj"] original_name = data[-1]["label"] self.assertEqual(reverted_name, original_name) print("invert diff", reverted.size - n_good) # 25300: 2 workers (cpu, non-macos) # 1812: 0 workers (gpu or macos) # 1821: windows torch 1.10.0 - self.assertTrue((reverted.size - n_good) in (34007, 1812, 1821), f"diff. {reverted.size - n_good}") + self.assertTrue((reverted.size - n_good) < 40000, f"diff. {reverted.size - n_good}") set_determinism(seed=None) diff --git a/tests/test_k_space_spike_noise.py b/tests/test_k_space_spike_noise.py index 85e8fec3c3..537655b3e5 100644 --- a/tests/test_k_space_spike_noise.py +++ b/tests/test_k_space_spike_noise.py @@ -53,9 +53,7 @@ def test_same_result(self, im_shape, im_type, k_intensity): out1 = t(deepcopy(im)) out2 = t(deepcopy(im)) - self.assertEqual(type(im), type(out1)) if isinstance(out1, torch.Tensor): - self.assertEqual(im.device, out1.device) out1 = out1.cpu() out2 = out2.cpu() @@ -69,9 +67,7 @@ def test_highlighted_kspace_pixel(self, im_shape, as_tensor_input, k_intensity): t = KSpaceSpikeNoise(loc, k_intensity) out = t(im) - self.assertEqual(type(im), type(out)) if isinstance(out, torch.Tensor): - self.assertEqual(im.device, out.device) out = out.cpu() if k_intensity is not None: diff --git a/tests/test_k_space_spike_noised.py b/tests/test_k_space_spike_noised.py index 0230f40b15..5a0aed7d67 100644 --- a/tests/test_k_space_spike_noised.py +++ b/tests/test_k_space_spike_noised.py @@ -57,9 +57,7 @@ def test_same_result(self, im_shape, im_type): out2 = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out1[k]), type(data[k])) if isinstance(out1[k], torch.Tensor): - self.assertEqual(out1[k].device, data[k].device) out1[k] = out1[k].cpu() out2[k] = out2[k].cpu() np.testing.assert_allclose(out1[k], out2[k]) @@ -75,9 +73,7 @@ def test_highlighted_kspace_pixel(self, im_shape, im_type): out = t(data) for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) out[k] = out[k].cpu() n_dims = len(im_shape) @@ -95,13 +91,6 @@ def test_dict_matches(self, im_shape, im_type): t = KSpaceSpikeNoised(KEYS, loc, k_intensity) out = t(deepcopy(data)) - - for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) - if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) - out[k] = out[k].cpu() - np.testing.assert_allclose(out[KEYS[0]], out[KEYS[1]]) diff --git a/tests/test_keep_largest_connected_component.py b/tests/test_keep_largest_connected_component.py index 5c96b62131..80dbc1c51d 100644 --- a/tests/test_keep_largest_connected_component.py +++ b/tests/test_keep_largest_connected_component.py @@ -19,7 +19,7 @@ from monai.transforms import KeepLargestConnectedComponent from monai.transforms.utils_pytorch_numpy_unification import moveaxis from monai.utils.type_conversion import convert_to_dst_type -from tests.utils import TEST_NDARRAYS, SkipIfBeforePyTorchVersion, assert_allclose +from tests.utils import TEST_NDARRAYS, assert_allclose def to_onehot(x): @@ -350,10 +350,9 @@ class TestKeepLargestConnectedComponent(unittest.TestCase): def test_correct_results(self, _, args, input_image, expected): converter = KeepLargestConnectedComponent(**args) result = converter(input_image) - assert_allclose(result, expected, type_test=False) + assert_allclose(result, expected, type_test="tensor") @parameterized.expand(TESTS) - @SkipIfBeforePyTorchVersion((1, 7)) def test_correct_results_before_after_onehot(self, _, args, input_image, expected): """ From torch==1.7, torch.argmax changes its mechanism that if there are multiple maximal values then the @@ -373,12 +372,12 @@ def test_correct_results_before_after_onehot(self, _, args, input_image, expecte img = to_onehot(input_image) result2 = KeepLargestConnectedComponent(**args)(img) result2 = result2.argmax(0)[None] - assert_allclose(result, result2) + assert_allclose(result, result2, type_test="tensor") # if onehotted, un-onehot and check result stays the same else: img = input_image.argmax(0)[None] result2 = KeepLargestConnectedComponent(**args)(img) - assert_allclose(result.argmax(0)[None], result2) + assert_allclose(result.argmax(0)[None], result2, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_keep_largest_connected_componentd.py b/tests/test_keep_largest_connected_componentd.py index a06fb51a97..544b7f1773 100644 --- a/tests/test_keep_largest_connected_componentd.py +++ b/tests/test_keep_largest_connected_componentd.py @@ -339,7 +339,7 @@ class TestKeepLargestConnectedComponentd(unittest.TestCase): def test_correct_results(self, _, args, input_dict, expected): converter = KeepLargestConnectedComponentd(**args) result = converter(input_dict) - assert_allclose(result["img"], expected, type_test=False) + assert_allclose(result["img"], expected, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_label_to_contour.py b/tests/test_label_to_contour.py index fef40af08d..cab116afbe 100644 --- a/tests/test_label_to_contour.py +++ b/tests/test_label_to_contour.py @@ -152,7 +152,7 @@ def test_contour(self): channels = cube.shape[0] for channel in range(channels): - assert_allclose(test_result_cube[channel, ...], expected_output) + assert_allclose(test_result_cube[channel, ...], expected_output, type_test="tensor") # check 4-dim input data test_img, expected_output = gen_fixed_img(p) @@ -162,7 +162,7 @@ def test_contour(self): self.assertEqual(test_result_img.shape, img.shape) for channel in range(channels): - assert_allclose(test_result_img[channel, ...], expected_output) + assert_allclose(test_result_img[channel, ...], expected_output, type_test="tensor") # check invalid input data error_input = torch.rand(1, 2) diff --git a/tests/test_label_to_contourd.py b/tests/test_label_to_contourd.py index 6481e803ba..e11b59130a 100644 --- a/tests/test_label_to_contourd.py +++ b/tests/test_label_to_contourd.py @@ -154,7 +154,7 @@ def test_contour(self): test_result_np = test_result_cube["img"] channels = cube.shape[0] for channel in range(channels): - assert_allclose(test_result_np[channel, ...], expected_output) + assert_allclose(test_result_np[channel, ...], expected_output, type_test="tensor") # check 4-dim input data test_img, expected_output = gen_fixed_img(p) @@ -165,7 +165,7 @@ def test_contour(self): test_result_np = test_result_img["img"] for channel in range(channels): - assert_allclose(test_result_np[channel, ...], expected_output) + assert_allclose(test_result_np[channel, ...], expected_output, type_test="tensor") # check invalid input data error_input = {"img": torch.rand(1, 2)} diff --git a/tests/test_label_to_mask.py b/tests/test_label_to_mask.py index 8f81a8da1a..cf25a4024a 100644 --- a/tests/test_label_to_mask.py +++ b/tests/test_label_to_mask.py @@ -12,7 +12,6 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import LabelToMask @@ -61,10 +60,7 @@ class TestLabelToMask(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, image, expected_data): result = LabelToMask(**argments)(image) - self.assertEqual(type(result), type(image)) - if isinstance(result, torch.Tensor): - self.assertEqual(result.device, image.device) - assert_allclose(result, expected_data, type_test=False) + assert_allclose(result, expected_data, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_label_to_maskd.py b/tests/test_label_to_maskd.py index e67b857502..5b194bd632 100644 --- a/tests/test_label_to_maskd.py +++ b/tests/test_label_to_maskd.py @@ -12,7 +12,6 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import LabelToMaskd @@ -61,11 +60,8 @@ class TestLabelToMaskd(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, argments, input_data, expected_data): result = LabelToMaskd(**argments)(input_data) - r, i = result["img"], input_data["img"] - self.assertEqual(type(r), type(i)) - if isinstance(r, torch.Tensor): - self.assertEqual(r.device, i.device) - assert_allclose(r, expected_data, type_test=False) + r = result["img"] + assert_allclose(r, expected_data, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_lambda.py b/tests/test_lambda.py index c187cc979b..3fa080f794 100644 --- a/tests/test_lambda.py +++ b/tests/test_lambda.py @@ -11,6 +11,7 @@ import unittest +from monai.data.meta_tensor import MetaTensor from monai.transforms.utility.array import Lambda from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose @@ -24,7 +25,7 @@ def identity_func(x): return x lambd = Lambda(func=identity_func) - assert_allclose(identity_func(img), lambd(img)) + assert_allclose(identity_func(img), lambd(img), type_test=False) def test_lambda_slicing(self): for p in TEST_NDARRAYS: @@ -34,7 +35,12 @@ def slice_func(x): return x[:, :, :6, ::2] lambd = Lambda(func=slice_func) - assert_allclose(slice_func(img), lambd(img)) + out = lambd(img) + assert_allclose(slice_func(img), out, type_test=False) + self.assertIsInstance(out, MetaTensor) + self.assertEqual(len(out.applied_operations), 1) + out = lambd.inverse(out) + self.assertEqual(len(out.applied_operations), 0) if __name__ == "__main__": diff --git a/tests/test_lambdad.py b/tests/test_lambdad.py index 30d70f40fb..0a582425e1 100644 --- a/tests/test_lambdad.py +++ b/tests/test_lambdad.py @@ -11,6 +11,7 @@ import unittest +from monai.data.meta_tensor import MetaTensor from monai.transforms.utility.dictionary import Lambdad from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose @@ -26,8 +27,8 @@ def noise_func(x): expected = {"img": noise_func(data["img"]), "prop": 1.0} ret = Lambdad(keys=["img", "prop"], func=noise_func, overwrite=[True, False])(data) - assert_allclose(expected["img"], ret["img"]) - assert_allclose(expected["prop"], ret["prop"]) + assert_allclose(expected["img"], ret["img"], type_test=False) + assert_allclose(expected["prop"], ret["prop"], type_test=False) def test_lambdad_slicing(self): for p in TEST_NDARRAYS: @@ -39,8 +40,15 @@ def slice_func(x): lambd = Lambdad(keys=data.keys(), func=slice_func) expected = {} - expected["img"] = slice_func(data["img"]) - assert_allclose(expected["img"], lambd(data)["img"]) + expected = slice_func(data["img"]) + out = lambd(data) + out_img = out["img"] + assert_allclose(expected, out_img, type_test=False) + self.assertIsInstance(out_img, MetaTensor) + self.assertEqual(len(out_img.applied_operations), 1) + inv_img = lambd.inverse(out)["img"] + self.assertIsInstance(inv_img, MetaTensor) + self.assertEqual(len(inv_img.applied_operations), 0) if __name__ == "__main__": diff --git a/tests/test_lesion_froc.py b/tests/test_lesion_froc.py index 6b4989a9d5..b135b3eaeb 100644 --- a/tests/test_lesion_froc.py +++ b/tests/test_lesion_froc.py @@ -31,7 +31,7 @@ def save_as_tif(filename, array): if not filename.endswith(".tif"): filename += ".tif" file_path = os.path.join("tests", "testing_data", filename) - imwrite(file_path, array, compress="jpeg", tile=(16, 16)) + imwrite(file_path, array, compression="jpeg", tile=(16, 16)) def around(val, interval=3): diff --git a/tests/test_load_decathlon_datalist.py b/tests/test_load_decathlon_datalist.py index 91d144d84f..a58ba73ece 100644 --- a/tests/test_load_decathlon_datalist.py +++ b/tests/test_load_decathlon_datalist.py @@ -29,7 +29,7 @@ def test_seg_values(self): {"image": "spleen_19.nii.gz", "label": "spleen_19.nii.gz"}, {"image": "spleen_31.nii.gz", "label": "spleen_31.nii.gz"}, ], - "test": ["spleen_15.nii.gz", "spleen_23.nii.gz"], + "test": [{"image": "spleen_15.nii.gz"}, {"image": "spleen_23.nii.gz"}], } json_str = json.dumps(test_data) file_path = os.path.join(tempdir, "test_data.json") @@ -38,6 +38,8 @@ def test_seg_values(self): result = load_decathlon_datalist(file_path, True, "training", tempdir) self.assertEqual(result[0]["image"], os.path.join(tempdir, "spleen_19.nii.gz")) self.assertEqual(result[0]["label"], os.path.join(tempdir, "spleen_19.nii.gz")) + result = load_decathlon_datalist(file_path, True, "test", None) + self.assertEqual(result[0]["image"], os.path.join(tempdir, "spleen_15.nii.gz")) def test_cls_values(self): with tempfile.TemporaryDirectory() as tempdir: @@ -81,6 +83,8 @@ def test_seg_no_basedir(self): result = load_decathlon_datalist(file_path, True, "training", None) self.assertEqual(result[0]["image"], os.path.join(tempdir, "spleen_19.nii.gz")) self.assertEqual(result[0]["label"], os.path.join(tempdir, "spleen_19.nii.gz")) + result = load_decathlon_datalist(file_path, True, "test", None) + self.assertEqual(result[0]["image"], os.path.join(tempdir, "spleen_15.nii.gz")) def test_seg_no_labels(self): with tempfile.TemporaryDirectory() as tempdir: diff --git a/tests/test_load_image.py b/tests/test_load_image.py index 201fe2fd5b..aad82df1db 100644 --- a/tests/test_load_image.py +++ b/tests/test_load_image.py @@ -10,6 +10,7 @@ # limitations under the License. import os +import shutil import tempfile import unittest from pathlib import Path @@ -17,11 +18,15 @@ import itk import nibabel as nib import numpy as np +import torch from parameterized import parameterized from PIL import Image -from monai.data import ITKReader, NibabelReader +from monai.data import ITKReader, NibabelReader, PydicomReader +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.transforms import LoadImage +from tests.utils import assert_allclose class _MiniReader: @@ -40,75 +45,57 @@ def get_data(self, _obj): return np.zeros((1, 1, 1)), {"name": "my test"} -TEST_CASE_1 = [{"image_only": True}, ["test_image.nii.gz"], (128, 128, 128)] +TEST_CASE_1 = [{}, ["test_image.nii.gz"], (128, 128, 128)] -TEST_CASE_2 = [{"image_only": False}, ["test_image.nii.gz"], (128, 128, 128)] +TEST_CASE_2 = [{}, ["test_image.nii.gz"], (128, 128, 128)] -TEST_CASE_3 = [ - {"image_only": True}, - ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], - (3, 128, 128, 128), -] +TEST_CASE_3 = [{}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (3, 128, 128, 128)] TEST_CASE_3_1 = [ # .mgz format - {"image_only": True, "reader": "nibabelreader"}, + {"reader": "nibabelreader"}, ["test_image.mgz", "test_image2.mgz", "test_image3.mgz"], (3, 128, 128, 128), ] -TEST_CASE_4 = [ - {"image_only": False}, - ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], - (3, 128, 128, 128), -] +TEST_CASE_4 = [{}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (3, 128, 128, 128)] TEST_CASE_4_1 = [ # additional parameter - {"image_only": False, "mmap": False}, + {"mmap": False}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (3, 128, 128, 128), ] -TEST_CASE_5 = [{"reader": NibabelReader(mmap=False), "image_only": False}, ["test_image.nii.gz"], (128, 128, 128)] +TEST_CASE_5 = [{"reader": NibabelReader(mmap=False)}, ["test_image.nii.gz"], (128, 128, 128)] -TEST_CASE_6 = [{"reader": ITKReader(), "image_only": True}, ["test_image.nii.gz"], (128, 128, 128)] +TEST_CASE_6 = [{"reader": ITKReader()}, ["test_image.nii.gz"], (128, 128, 128)] -TEST_CASE_7 = [{"reader": ITKReader(), "image_only": False}, ["test_image.nii.gz"], (128, 128, 128)] +TEST_CASE_7 = [{"reader": ITKReader()}, ["test_image.nii.gz"], (128, 128, 128)] TEST_CASE_8 = [ - {"reader": ITKReader(), "image_only": True}, + {"reader": ITKReader()}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (3, 128, 128, 128), ] TEST_CASE_8_1 = [ - {"reader": ITKReader(channel_dim=0), "image_only": True}, + {"reader": ITKReader(channel_dim=0)}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (384, 128, 128), ] TEST_CASE_9 = [ - {"reader": ITKReader(), "image_only": False}, + {"reader": ITKReader()}, ["test_image.nii.gz", "test_image2.nii.gz", "test_image3.nii.gz"], (3, 128, 128, 128), ] -TEST_CASE_10 = [ - {"image_only": False, "reader": ITKReader(pixel_type=itk.UC)}, - "tests/testing_data/CT_DICOM", - (16, 16, 4), - (16, 16, 4), -] +TEST_CASE_10 = [{"reader": ITKReader(pixel_type=itk.UC)}, "tests/testing_data/CT_DICOM", (16, 16, 4), (16, 16, 4)] -TEST_CASE_11 = [ - {"image_only": False, "reader": "ITKReader", "pixel_type": itk.UC}, - "tests/testing_data/CT_DICOM", - (16, 16, 4), - (16, 16, 4), -] +TEST_CASE_11 = [{"reader": "ITKReader", "pixel_type": itk.UC}, "tests/testing_data/CT_DICOM", (16, 16, 4), (16, 16, 4)] TEST_CASE_12 = [ - {"image_only": False, "reader": "ITKReader", "pixel_type": itk.UC, "reverse_indexing": True}, + {"reader": "ITKReader", "pixel_type": itk.UC, "reverse_indexing": True}, "tests/testing_data/CT_DICOM", (16, 16, 4), (4, 16, 16), @@ -134,6 +121,32 @@ def get_data(self, _obj): (128, 128, 3, 128), ] +# test same dicom data with PydicomReader +TEST_CASE_19 = [{"reader": PydicomReader()}, "tests/testing_data/CT_DICOM", (16, 16, 4), (16, 16, 4)] + +TEST_CASE_20 = [ + {"reader": "PydicomReader", "ensure_channel_first": True}, + "tests/testing_data/CT_DICOM", + (16, 16, 4), + (1, 16, 16, 4), +] + +TEST_CASE_21 = [ + {"reader": "PydicomReader", "affine_lps_to_ras": True, "defer_size": "2 MB"}, + "tests/testing_data/CT_DICOM", + (16, 16, 4), + (16, 16, 4), +] + +# test reader consistency between PydicomReader and ITKReader on dicom data +TEST_CASE_22 = ["tests/testing_data/CT_DICOM"] + + +TESTS_META = [] +for track_meta in (False, True): + TESTS_META.append([{}, (128, 128, 128), track_meta]) + TESTS_META.append([{"reader": "ITKReader", "fallback_only": False}, (128, 128, 128), track_meta]) + class TestLoadImage(unittest.TestCase): @parameterized.expand( @@ -145,14 +158,10 @@ def test_nibabel_reader(self, input_param, filenames, expected_shape): for i, name in enumerate(filenames): filenames[i] = os.path.join(tempdir, name) nib.save(nib.Nifti1Image(test_image, np.eye(4)), filenames[i]) - result = LoadImage(**input_param)(filenames) - - if isinstance(result, tuple): - result, header = result - self.assertTrue("affine" in header) - self.assertEqual(header["filename_or_obj"], os.path.join(tempdir, "test_image.nii.gz")) - np.testing.assert_allclose(header["affine"], np.eye(4)) - np.testing.assert_allclose(header["original_affine"], np.eye(4)) + result = LoadImage(image_only=True, **input_param)(filenames) + ext = "".join(Path(name).suffixes) + self.assertEqual(result.meta["filename_or_obj"], os.path.join(tempdir, "test_image" + ext)) + assert_allclose(result.affine, torch.eye(4)) self.assertTupleEqual(result.shape, expected_shape) @parameterized.expand([TEST_CASE_6, TEST_CASE_7, TEST_CASE_8, TEST_CASE_8_1, TEST_CASE_9]) @@ -163,25 +172,19 @@ def test_itk_reader(self, input_param, filenames, expected_shape): filenames[i] = os.path.join(tempdir, name) itk_np_view = itk.image_view_from_array(test_image) itk.imwrite(itk_np_view, filenames[i]) - result = LoadImage(**input_param)(filenames) - - if isinstance(result, tuple): - result, header = result - self.assertTrue("affine" in header) - self.assertEqual(header["filename_or_obj"], os.path.join(tempdir, "test_image.nii.gz")) - np_diag = np.diag([-1, -1, 1, 1]) - np.testing.assert_allclose(header["affine"], np_diag) - np.testing.assert_allclose(header["original_affine"], np_diag) + result = LoadImage(image_only=True, **input_param)(filenames) + self.assertEqual(result.meta["filename_or_obj"], os.path.join(tempdir, "test_image.nii.gz")) + diag = torch.as_tensor(np.diag([-1, -1, 1, 1])) + np.testing.assert_allclose(result.affine, diag) self.assertTupleEqual(result.shape, expected_shape) - @parameterized.expand([TEST_CASE_10, TEST_CASE_11, TEST_CASE_12]) + @parameterized.expand([TEST_CASE_10, TEST_CASE_11, TEST_CASE_12, TEST_CASE_19, TEST_CASE_20, TEST_CASE_21]) def test_itk_dicom_series_reader(self, input_param, filenames, expected_shape, expected_np_shape): - result, header = LoadImage(**input_param)(filenames) - self.assertTrue("affine" in header) - self.assertEqual(header["filename_or_obj"], f"{Path(filenames)}") - np.testing.assert_allclose( - header["affine"], - np.array( + result = LoadImage(image_only=True, **input_param)(filenames) + self.assertEqual(result.meta["filename_or_obj"], f"{Path(filenames)}") + assert_allclose( + result.affine, + torch.tensor( [ [-0.488281, 0.0, 0.0, 125.0], [0.0, -0.488281, 0.0, 128.100006], @@ -190,7 +193,6 @@ def test_itk_dicom_series_reader(self, input_param, filenames, expected_shape, e ] ), ) - self.assertTupleEqual(tuple(header["spatial_shape"]), expected_shape) self.assertTupleEqual(result.shape, expected_np_shape) def test_itk_reader_multichannel(self): @@ -200,14 +202,24 @@ def test_itk_reader_multichannel(self): itk_np_view = itk.image_view_from_array(test_image, is_vector=True) itk.imwrite(itk_np_view, filename) for flag in (False, True): - result, header = LoadImage(reader=ITKReader(reverse_indexing=flag))(Path(filename)) - - self.assertTupleEqual(tuple(header["spatial_shape"]), (224, 256)) + result = LoadImage(image_only=True, reader=ITKReader(reverse_indexing=flag))(Path(filename)) test_image = test_image.transpose(1, 0, 2) np.testing.assert_allclose(result[:, :, 0], test_image[:, :, 0]) np.testing.assert_allclose(result[:, :, 1], test_image[:, :, 1]) np.testing.assert_allclose(result[:, :, 2], test_image[:, :, 2]) + @parameterized.expand([TEST_CASE_22]) + def test_dicom_reader_consistency(self, filenames): + itk_param = {"reader": "ITKReader"} + pydicom_param = {"reader": "PydicomReader"} + for affine_flag in [True, False]: + itk_param["affine_lps_to_ras"] = affine_flag + pydicom_param["affine_lps_to_ras"] = affine_flag + itk_result = LoadImage(image_only=True, **itk_param)(filenames) + pydicom_result = LoadImage(image_only=True, **pydicom_param)(filenames) + np.testing.assert_allclose(pydicom_result, itk_result) + np.testing.assert_allclose(pydicom_result.affine, itk_result.affine) + def test_load_nifti_multichannel(self): test_image = np.random.randint(0, 256, size=(31, 64, 16, 2)).astype(np.float32) with tempfile.TemporaryDirectory() as tempdir: @@ -215,12 +227,10 @@ def test_load_nifti_multichannel(self): itk_np_view = itk.image_view_from_array(test_image, is_vector=True) itk.imwrite(itk_np_view, filename) - itk_img, itk_header = LoadImage(reader=ITKReader())(Path(filename)) - self.assertTupleEqual(tuple(itk_header["spatial_shape"]), (16, 64, 31)) + itk_img = LoadImage(image_only=True, reader=ITKReader())(Path(filename)) self.assertTupleEqual(tuple(itk_img.shape), (16, 64, 31, 2)) - nib_image, nib_header = LoadImage(reader=NibabelReader(squeeze_non_spatial_dims=True))(Path(filename)) - self.assertTupleEqual(tuple(nib_header["spatial_shape"]), (16, 64, 31)) + nib_image = LoadImage(image_only=True, reader=NibabelReader(squeeze_non_spatial_dims=True))(Path(filename)) self.assertTupleEqual(tuple(nib_image.shape), (16, 64, 31, 2)) np.testing.assert_allclose(itk_img, nib_image, atol=1e-3, rtol=1e-3) @@ -231,8 +241,7 @@ def test_load_png(self): with tempfile.TemporaryDirectory() as tempdir: filename = os.path.join(tempdir, "test_image.png") Image.fromarray(test_image.astype("uint8")).save(filename) - result, header = LoadImage(image_only=False)(filename) - self.assertTupleEqual(tuple(header["spatial_shape"]), spatial_size[::-1]) + result = LoadImage(image_only=True)(filename) self.assertTupleEqual(result.shape, spatial_size[::-1]) np.testing.assert_allclose(result.T, test_image) @@ -244,10 +253,9 @@ def test_register(self): itk_np_view = itk.image_view_from_array(test_image) itk.imwrite(itk_np_view, filename) - loader = LoadImage(image_only=False) + loader = LoadImage(image_only=True) loader.register(ITKReader()) - result, header = loader(filename) - self.assertTupleEqual(tuple(header["spatial_shape"]), spatial_size[::-1]) + result = loader(filename) self.assertTupleEqual(result.shape, spatial_size[::-1]) def test_kwargs(self): @@ -258,35 +266,37 @@ def test_kwargs(self): itk_np_view = itk.image_view_from_array(test_image) itk.imwrite(itk_np_view, filename) - loader = LoadImage(image_only=False) + loader = LoadImage(image_only=True) reader = ITKReader(fallback_only=False) loader.register(reader) - result, header = loader(filename) + result = loader(filename) reader = ITKReader() img = reader.read(filename, fallback_only=False) - result_raw, header_raw = reader.get_data(img) - np.testing.assert_allclose(header["spatial_shape"], header_raw["spatial_shape"]) + result_raw = reader.get_data(img) + result_raw = MetaTensor.ensure_torch_and_prune_meta(*result_raw) self.assertTupleEqual(result.shape, result_raw.shape) def test_my_reader(self): """test customised readers""" - out = LoadImage(reader=_MiniReader, is_compatible=True)("test") - self.assertEqual(out[1]["name"], "my test") - out = LoadImage(reader=_MiniReader, is_compatible=False)("test") - self.assertEqual(out[1]["name"], "my test") + out = LoadImage(image_only=True, reader=_MiniReader, is_compatible=True)("test") + self.assertEqual(out.meta["name"], "my test") + out = LoadImage(image_only=True, reader=_MiniReader, is_compatible=False)("test") + self.assertEqual(out.meta["name"], "my test") for item in (_MiniReader, _MiniReader(is_compatible=False)): - out = LoadImage(reader=item)("test") - self.assertEqual(out[1]["name"], "my test") - out = LoadImage()("test", reader=_MiniReader(is_compatible=False)) - self.assertEqual(out[1]["name"], "my test") + out = LoadImage(image_only=True, reader=item)("test") + self.assertEqual(out.meta["name"], "my test") + out = LoadImage(image_only=True)("test", reader=_MiniReader(is_compatible=False)) + self.assertEqual(out.meta["name"], "my test") def test_itk_meta(self): """test metadata from a directory""" - out, meta = LoadImage(reader="ITKReader", pixel_type=itk.UC, series_meta=True)("tests/testing_data/CT_DICOM") + out = LoadImage(image_only=True, reader="ITKReader", pixel_type=itk.UC, series_meta=True)( + "tests/testing_data/CT_DICOM" + ) idx = "0008|103e" label = itk.GDCMImageIO.GetLabelFromTag(idx, "")[1] - val = meta[idx] + val = out.meta[idx] expected = "Series Description=Routine Brain " self.assertEqual(f"{label}={val}", expected) @@ -296,13 +306,41 @@ def test_channel_dim(self, input_param, filename, expected_shape): with tempfile.TemporaryDirectory() as tempdir: filename = os.path.join(tempdir, filename) nib.save(nib.Nifti1Image(test_image, np.eye(4)), filename) - result = LoadImage(**input_param)(filename) + result = LoadImage(image_only=True, **input_param)(filename) self.assertTupleEqual( - result[0].shape, (3, 128, 128, 128) if input_param.get("ensure_channel_first", False) else expected_shape + result.shape, (3, 128, 128, 128) if input_param.get("ensure_channel_first", False) else expected_shape ) - self.assertTupleEqual(tuple(result[1]["spatial_shape"]), (128, 128, 128)) - self.assertEqual(result[1]["original_channel_dim"], input_param["channel_dim"]) + self.assertEqual(result.meta["original_channel_dim"], input_param["channel_dim"]) + + +class TestLoadImageMeta(unittest.TestCase): + @classmethod + def setUpClass(cls): + super(__class__, cls).setUpClass() + cls.tmpdir = tempfile.mkdtemp() + test_image = nib.Nifti1Image(np.random.rand(128, 128, 128), np.eye(4)) + nib.save(test_image, os.path.join(cls.tmpdir, "im.nii.gz")) + cls.test_data = os.path.join(cls.tmpdir, "im.nii.gz") + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tmpdir) + super(__class__, cls).tearDownClass() + + @parameterized.expand(TESTS_META) + def test_correct(self, input_param, expected_shape, track_meta): + set_track_meta(track_meta) + r = LoadImage(image_only=True, **input_param)(self.test_data) + self.assertTupleEqual(r.shape, expected_shape) + if track_meta: + self.assertIsInstance(r, MetaTensor) + self.assertTrue(hasattr(r, "affine")) + self.assertIsInstance(r.affine, torch.Tensor) + else: + self.assertIsInstance(r, torch.Tensor) + self.assertNotIsInstance(r, MetaTensor) + self.assertFalse(hasattr(r, "affine")) if __name__ == "__main__": diff --git a/tests/test_load_imaged.py b/tests/test_load_imaged.py index bc001cf2fd..1f2ce31ad2 100644 --- a/tests/test_load_imaged.py +++ b/tests/test_load_imaged.py @@ -10,6 +10,7 @@ # limitations under the License. import os +import shutil import tempfile import unittest from pathlib import Path @@ -17,11 +18,15 @@ import itk import nibabel as nib import numpy as np +import torch from parameterized import parameterized from monai.data import ITKReader -from monai.transforms import Compose, EnsureChannelFirstD, LoadImaged, SaveImageD -from monai.utils.enums import PostFix +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor +from monai.transforms import Compose, EnsureChannelFirstD, FromMetaTensord, LoadImaged, SaveImageD +from monai.transforms.meta_utility.dictionary import ToMetaTensord +from tests.utils import assert_allclose KEYS = ["image", "label", "extra"] @@ -29,6 +34,11 @@ TEST_CASE_2 = [{"keys": KEYS, "reader": "ITKReader", "fallback_only": False}, (128, 128, 128)] +TESTS_META = [] +for track_meta in (False, True): + TESTS_META.append([{"keys": KEYS}, (128, 128, 128), track_meta]) + TESTS_META.append([{"keys": KEYS, "reader": "ITKReader", "fallback_only": False}, (128, 128, 128), track_meta]) + class TestLoadImaged(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) @@ -39,7 +49,7 @@ def test_shape(self, input_param, expected_shape): for key in KEYS: nib.save(test_image, os.path.join(tempdir, key + ".nii.gz")) test_data.update({key: os.path.join(tempdir, key + ".nii.gz")}) - result = LoadImaged(**input_param)(test_data) + result = LoadImaged(image_only=True, **input_param)(test_data) for key in KEYS: self.assertTupleEqual(result[key].shape, expected_shape) @@ -52,10 +62,9 @@ def test_register(self): itk_np_view = itk.image_view_from_array(test_image) itk.imwrite(itk_np_view, filename) - loader = LoadImaged(keys="img") + loader = LoadImaged(keys="img", image_only=True) loader.register(ITKReader()) result = loader({"img": Path(filename)}) - self.assertTupleEqual(tuple(result[PostFix.meta("img")]["spatial_shape"]), spatial_size[::-1]) self.assertTupleEqual(result["img"].shape, spatial_size[::-1]) def test_channel_dim(self): @@ -65,63 +74,64 @@ def test_channel_dim(self): filename = os.path.join(tempdir, "test_image.nii.gz") nib.save(nib.Nifti1Image(test_image, affine=np.eye(4)), filename) - loader = LoadImaged(keys="img") + loader = LoadImaged(keys="img", image_only=True) loader.register(ITKReader(channel_dim=2)) - result = EnsureChannelFirstD("img")(loader({"img": filename})) - self.assertTupleEqual(tuple(result[PostFix.meta("img")]["spatial_shape"]), (32, 64, 128)) + t = Compose([EnsureChannelFirstD("img"), FromMetaTensord("img")]) + result = t(loader({"img": filename})) self.assertTupleEqual(result["img"].shape, (3, 32, 64, 128)) def test_no_file(self): with self.assertRaises(RuntimeError): - LoadImaged(keys="img")({"img": "unknown"}) + LoadImaged(keys="img", image_only=True)({"img": "unknown"}) with self.assertRaises(RuntimeError): - LoadImaged(keys="img", reader="nibabelreader")({"img": "unknown"}) + LoadImaged(keys="img", reader="nibabelreader", image_only=True)({"img": "unknown"}) class TestConsistency(unittest.TestCase): - def _cmp(self, filename, shape, ch_shape, reader_1, reader_2, outname, ext): + def _cmp(self, filename, ch_shape, reader_1, reader_2, outname, ext): data_dict = {"img": filename} keys = data_dict.keys() - xforms = Compose([LoadImaged(keys, reader=reader_1, ensure_channel_first=True)]) + xforms = Compose([LoadImaged(keys, reader=reader_1, ensure_channel_first=True, image_only=True)]) img_dict = xforms(data_dict) # load dicom with itk self.assertTupleEqual(img_dict["img"].shape, ch_shape) - self.assertTupleEqual(tuple(img_dict[PostFix.meta("img")]["spatial_shape"]), shape) with tempfile.TemporaryDirectory() as tempdir: - save_xform = SaveImageD( - keys, meta_keys=PostFix.meta("img"), output_dir=tempdir, squeeze_end_dims=False, output_ext=ext - ) + save_xform = SaveImageD(keys, output_dir=tempdir, squeeze_end_dims=False, output_ext=ext) save_xform(img_dict) # save to nifti - new_xforms = Compose([LoadImaged(keys, reader=reader_2), EnsureChannelFirstD(keys)]) + new_xforms = Compose( + [ + LoadImaged(keys, reader=reader_2, image_only=True), + EnsureChannelFirstD(keys), + FromMetaTensord(keys), + ToMetaTensord(keys), + ] + ) out = new_xforms({"img": os.path.join(tempdir, outname)}) # load nifti with itk self.assertTupleEqual(out["img"].shape, ch_shape) - self.assertTupleEqual(tuple(out[PostFix.meta("img")]["spatial_shape"]), shape) - if "affine" in img_dict[PostFix.meta("img")] and "affine" in out[PostFix.meta("img")]: - np.testing.assert_allclose( - img_dict[PostFix.meta("img")]["affine"], out[PostFix.meta("img")]["affine"], rtol=1e-3 - ) - np.testing.assert_allclose(out["img"], img_dict["img"], rtol=1e-3) + + def is_identity(x): + return (x == torch.eye(x.shape[0])).all() + + if not is_identity(img_dict["img"].affine) and not is_identity(out["img"].affine): + assert_allclose(img_dict["img"].affine, out["img"].affine, rtol=1e-3) + assert_allclose(out["img"], img_dict["img"], rtol=1e-3) def test_dicom(self): img_dir = "tests/testing_data/CT_DICOM" - self._cmp( - img_dir, (16, 16, 4), (1, 16, 16, 4), "itkreader", "itkreader", "CT_DICOM/CT_DICOM_trans.nii.gz", ".nii.gz" - ) + self._cmp(img_dir, (1, 16, 16, 4), "itkreader", "itkreader", "CT_DICOM/CT_DICOM_trans.nii.gz", ".nii.gz") output_name = "CT_DICOM/CT_DICOM_trans.nii.gz" - self._cmp(img_dir, (16, 16, 4), (1, 16, 16, 4), "nibabelreader", "itkreader", output_name, ".nii.gz") - self._cmp(img_dir, (16, 16, 4), (1, 16, 16, 4), "itkreader", "nibabelreader", output_name, ".nii.gz") + self._cmp(img_dir, (1, 16, 16, 4), "nibabelreader", "itkreader", output_name, ".nii.gz") + self._cmp(img_dir, (1, 16, 16, 4), "itkreader", "nibabelreader", output_name, ".nii.gz") def test_multi_dicom(self): """multichannel dicom reading, saving to nifti, then load with itk or nibabel""" img_dir = ["tests/testing_data/CT_DICOM", "tests/testing_data/CT_DICOM"] - self._cmp( - img_dir, (16, 16, 4), (2, 16, 16, 4), "itkreader", "itkreader", "CT_DICOM/CT_DICOM_trans.nii.gz", ".nii.gz" - ) + self._cmp(img_dir, (2, 16, 16, 4), "itkreader", "itkreader", "CT_DICOM/CT_DICOM_trans.nii.gz", ".nii.gz") output_name = "CT_DICOM/CT_DICOM_trans.nii.gz" - self._cmp(img_dir, (16, 16, 4), (2, 16, 16, 4), "nibabelreader", "itkreader", output_name, ".nii.gz") - self._cmp(img_dir, (16, 16, 4), (2, 16, 16, 4), "itkreader", "nibabelreader", output_name, ".nii.gz") + self._cmp(img_dir, (2, 16, 16, 4), "nibabelreader", "itkreader", output_name, ".nii.gz") + self._cmp(img_dir, (2, 16, 16, 4), "itkreader", "nibabelreader", output_name, ".nii.gz") def test_png(self): """png reading with itk, saving to nifti, then load with itk or nibabel or PIL""" @@ -132,9 +142,45 @@ def test_png(self): itk_np_view = itk.image_view_from_array(test_image, is_vector=True) itk.imwrite(itk_np_view, filename) output_name = "test_image/test_image_trans.png" - self._cmp(filename, (224, 256), (3, 224, 256), "itkreader", "itkreader", output_name, ".png") - self._cmp(filename, (224, 256), (3, 224, 256), "itkreader", "PILReader", output_name, ".png") - self._cmp(filename, (224, 256), (3, 224, 256), "itkreader", "nibabelreader", output_name, ".png") + self._cmp(filename, (3, 224, 256), "itkreader", "itkreader", output_name, ".png") + self._cmp(filename, (3, 224, 256), "itkreader", "PILReader", output_name, ".png") + self._cmp(filename, (3, 224, 256), "itkreader", "nibabelreader", output_name, ".png") + + +class TestLoadImagedMeta(unittest.TestCase): + @classmethod + def setUpClass(cls): + super(__class__, cls).setUpClass() + cls.tmpdir = tempfile.mkdtemp() + test_image = nib.Nifti1Image(np.random.rand(128, 128, 128), np.eye(4)) + cls.test_data = {} + for key in KEYS: + nib.save(test_image, os.path.join(cls.tmpdir, key + ".nii.gz")) + cls.test_data.update({key: os.path.join(cls.tmpdir, key + ".nii.gz")}) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tmpdir) + super(__class__, cls).tearDownClass() + + @parameterized.expand(TESTS_META) + def test_correct(self, input_param, expected_shape, track_meta): + set_track_meta(track_meta) + result = LoadImaged(image_only=True, **input_param)(self.test_data) + + # shouldn't have any extra meta data keys + self.assertEqual(len(result), len(KEYS)) + for key in KEYS: + r = result[key] + self.assertTupleEqual(r.shape, expected_shape) + if track_meta: + self.assertIsInstance(r, MetaTensor) + self.assertTrue(hasattr(r, "affine")) + self.assertIsInstance(r.affine, torch.Tensor) + else: + self.assertIsInstance(r, torch.Tensor) + self.assertNotIsInstance(r, MetaTensor) + self.assertFalse(hasattr(r, "affine")) if __name__ == "__main__": diff --git a/tests/test_load_spacing_orientation.py b/tests/test_load_spacing_orientation.py index 2792822c3d..8257b9965f 100644 --- a/tests/test_load_spacing_orientation.py +++ b/tests/test_load_spacing_orientation.py @@ -15,11 +15,11 @@ import nibabel import numpy as np +import torch from nibabel.processing import resample_to_output from parameterized import parameterized -from monai.transforms import AddChanneld, LoadImaged, Orientationd, Spacingd -from monai.utils.enums import PostFix +from monai.transforms import AddChanneld, Compose, LoadImaged, Orientationd, Spacingd FILES = tuple( os.path.join(os.path.dirname(__file__), "testing_data", filename) @@ -28,43 +28,45 @@ class TestLoadSpacingOrientation(unittest.TestCase): + @staticmethod + def load_image(filename): + data = {"image": filename} + t = Compose([LoadImaged(keys="image"), AddChanneld(keys="image")]) + return t(data) + @parameterized.expand(FILES) def test_load_spacingd(self, filename): - data = {"image": filename} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) + data_dict = self.load_image(filename) t = time.time() res_dict = Spacingd(keys="image", pixdim=(1, 0.2, 1), diagonal=True, padding_mode="zeros")(data_dict) t1 = time.time() print(f"time monai: {t1 - t}") - anat = nibabel.Nifti1Image(data_dict["image"][0], data_dict[PostFix.meta("image")]["original_affine"]) + anat = nibabel.Nifti1Image(np.asarray(data_dict["image"][0]), data_dict["image"].meta["original_affine"]) ref = resample_to_output(anat, (1, 0.2, 1), order=1) t2 = time.time() print(f"time scipy: {t2 - t1}") self.assertTrue(t2 >= t1) - np.testing.assert_allclose(res_dict[PostFix.meta("image")]["affine"], ref.affine) + np.testing.assert_allclose(res_dict["image"].affine, ref.affine) np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape) np.testing.assert_allclose(ref.get_fdata(), res_dict["image"][0], atol=0.05) @parameterized.expand(FILES) def test_load_spacingd_rotate(self, filename): - data = {"image": filename} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) - affine = data_dict[PostFix.meta("image")]["affine"] - data_dict[PostFix.meta("image")]["original_affine"] = data_dict[PostFix.meta("image")]["affine"] = ( - np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine + data_dict = self.load_image(filename) + affine = data_dict["image"].affine + data_dict["image"].meta["original_affine"] = data_dict["image"].affine = ( + torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]], dtype=torch.float64) @ affine ) t = time.time() res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=True, padding_mode="zeros")(data_dict) t1 = time.time() print(f"time monai: {t1 - t}") - anat = nibabel.Nifti1Image(data_dict["image"][0], data_dict[PostFix.meta("image")]["original_affine"]) + anat = nibabel.Nifti1Image(np.asarray(data_dict["image"][0]), data_dict["image"].meta["original_affine"]) ref = resample_to_output(anat, (1, 2, 3), order=1) t2 = time.time() print(f"time scipy: {t2 - t1}") self.assertTrue(t2 >= t1) - np.testing.assert_allclose(res_dict[PostFix.meta("image")]["affine"], ref.affine) + np.testing.assert_allclose(res_dict["image"].affine, ref.affine) if "anatomical" not in filename: np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape) np.testing.assert_allclose(ref.get_fdata(), res_dict["image"][0], atol=0.05) @@ -74,16 +76,14 @@ def test_load_spacingd_rotate(self, filename): np.testing.assert_allclose(ref.get_fdata()[..., :-1], res_dict["image"][0], atol=0.05) def test_load_spacingd_non_diag(self): - data = {"image": FILES[1]} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) - affine = data_dict[PostFix.meta("image")]["affine"] - data_dict[PostFix.meta("image")]["original_affine"] = data_dict[PostFix.meta("image")]["affine"] = ( - np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine + data_dict = self.load_image(FILES[1]) + affine = data_dict["image"].affine + data_dict["image"].meta["original_affine"] = data_dict["image"].affine = ( + torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]], dtype=torch.float64) @ affine ) res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="zeros")(data_dict) np.testing.assert_allclose( - res_dict[PostFix.meta("image")]["affine"], + res_dict["image"].affine, np.array( [ [0.0, 0.0, 3.0, -27.599409], @@ -95,38 +95,42 @@ def test_load_spacingd_non_diag(self): ) def test_load_spacingd_rotate_non_diag(self): - data = {"image": FILES[0]} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) + data_dict = self.load_image(FILES[0]) res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="border")(data_dict) np.testing.assert_allclose( - res_dict[PostFix.meta("image")]["affine"], + res_dict["image"].affine, np.array([[-1.0, 0.0, 0.0, 32.0], [0.0, 2.0, 0.0, -40.0], [0.0, 0.0, 3.0, -16.0], [0.0, 0.0, 0.0, 1.0]]), ) def test_load_spacingd_rotate_non_diag_ornt(self): - data = {"image": FILES[0]} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) - res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="border")(data_dict) - res_dict = Orientationd(keys="image", axcodes="LPI")(res_dict) + data_dict = self.load_image(FILES[0]) + t = Compose( + [ + Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="border"), + Orientationd(keys="image", axcodes="LPI"), + ] + ) + res_dict = t(data_dict) np.testing.assert_allclose( - res_dict[PostFix.meta("image")]["affine"], + res_dict["image"].affine, np.array([[-1.0, 0.0, 0.0, 32.0], [0.0, -2.0, 0.0, 40.0], [0.0, 0.0, -3.0, 32.0], [0.0, 0.0, 0.0, 1.0]]), ) def test_load_spacingd_non_diag_ornt(self): - data = {"image": FILES[1]} - data_dict = LoadImaged(keys="image")(data) - data_dict = AddChanneld(keys="image")(data_dict) - affine = data_dict[PostFix.meta("image")]["affine"] - data_dict[PostFix.meta("image")]["original_affine"] = data_dict[PostFix.meta("image")]["affine"] = ( - np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine + data_dict = self.load_image(FILES[1]) + affine = data_dict["image"].affine + data_dict["image"].meta["original_affine"] = data_dict["image"].affine = ( + torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]], dtype=torch.float64) @ affine ) - res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="border")(data_dict) - res_dict = Orientationd(keys="image", axcodes="LPI")(res_dict) + t = Compose( + [ + Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=False, padding_mode="border"), + Orientationd(keys="image", axcodes="LPI"), + ] + ) + res_dict = t(data_dict) np.testing.assert_allclose( - res_dict[PostFix.meta("image")]["affine"], + res_dict["image"].affine, np.array( [ [-3.0, 0.0, 0.0, 56.4005909], diff --git a/tests/test_look_up_option.py b/tests/test_look_up_option.py index 89fec1b575..700f7b9691 100644 --- a/tests/test_look_up_option.py +++ b/tests/test_look_up_option.py @@ -14,7 +14,7 @@ from parameterized import parameterized -from monai.utils import look_up_option +from monai.utils import StrEnum, look_up_option class _CaseEnum(Enum): @@ -27,6 +27,11 @@ class _CaseEnum1(Enum): EMPTY = "empty" +class _CaseStrEnum(StrEnum): + MODE_A = "A" + MODE_B = "B" + + TEST_CASES = ( ("test", ("test", "test1"), "test"), ("test1", {"test1", "test"}, "test1"), @@ -46,6 +51,14 @@ def test_default(self): output = look_up_option("not here", {"a", "b"}, default=None) self.assertEqual(output, None) + def test_str_enum(self): + output = look_up_option("C", {"A", "B"}, default=None) + self.assertEqual(output, None) + self.assertEqual(list(_CaseStrEnum), ["A", "B"]) + self.assertEqual(_CaseStrEnum.MODE_A, "A") + self.assertEqual(str(_CaseStrEnum.MODE_A), "A") + self.assertEqual(look_up_option("A", _CaseStrEnum), "A") + def test_no_found(self): with self.assertRaisesRegex(ValueError, "Unsupported"): look_up_option("not here", {"a", "b"}) diff --git a/tests/test_mask_intensity.py b/tests/test_mask_intensity.py index b6cfe0e10c..65182b6678 100644 --- a/tests/test_mask_intensity.py +++ b/tests/test_mask_intensity.py @@ -16,6 +16,7 @@ from parameterized import parameterized from monai.transforms import MaskIntensity +from tests.utils import TEST_NDARRAYS, assert_allclose TEST_CASE_1 = [ {"mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]]])}, @@ -54,8 +55,9 @@ class TestMaskIntensity(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5]) def test_value(self, argments, image, expected_data): - result = MaskIntensity(**argments)(image) - np.testing.assert_allclose(result, expected_data) + for p in TEST_NDARRAYS: + result = MaskIntensity(**argments)(p(image)) + assert_allclose(result, p(expected_data), type_test="tensor") def test_runtime_mask(self): mask_data = np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]]]) @@ -63,7 +65,7 @@ def test_runtime_mask(self): expected = np.array([[[0, 0, 0], [0, 2, 0], [0, 0, 0]], [[0, 0, 0], [0, 5, 0], [0, 0, 0]]]) result = MaskIntensity()(img=img, mask_data=mask_data) - np.testing.assert_allclose(result, expected) + assert_allclose(result, expected, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_masked_loss.py b/tests/test_masked_loss.py index 9f28d51aa4..a00b3ae7e7 100644 --- a/tests/test_masked_loss.py +++ b/tests/test_masked_loss.py @@ -17,7 +17,7 @@ from monai.losses.dice import DiceFocalLoss, DiceLoss from monai.losses.spatial_mask import MaskedLoss from monai.utils import set_determinism -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -73,7 +73,6 @@ def test_ill_opts(self): masked = MaskedLoss(loss=dice_loss) masked(input=torch.zeros((3, 3, 2, 2)), target=torch.zeros((3, 2, 2, 2)), mask=torch.zeros((3, 3, 2, 2))) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): input_param, expected_val = TEST_CASES[0] size = [3, 3, 5, 5] diff --git a/tests/test_masked_patch_wsi_dataset.py b/tests/test_masked_patch_wsi_dataset.py new file mode 100644 index 0000000000..a79bb5533b --- /dev/null +++ b/tests/test_masked_patch_wsi_dataset.py @@ -0,0 +1,93 @@ +# 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. + +import os +import unittest +from unittest import skipUnless + +from numpy.testing import assert_array_equal +from parameterized import parameterized + +from monai.data import MaskedPatchWSIDataset +from monai.utils import WSIPatchKeys, optional_import, set_determinism +from tests.utils import download_url_or_skip_test, testing_data_config + +set_determinism(0) + +cucim, has_cucim = optional_import("cucim") +has_cucim = has_cucim and hasattr(cucim, "CuImage") +_, has_osl = optional_import("openslide") +_, has_tiff = optional_import("tifffile", name="imwrite") +_, has_codec = optional_import("imagecodecs") +has_tiff = has_tiff and has_codec + + +FILE_KEY = "wsi_img" +FILE_URL = testing_data_config("images", FILE_KEY, "url") +base_name, extension = os.path.basename(f"{FILE_URL}"), ".tiff" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", "temp_" + base_name + extension) + +TEST_CASE_0 = [ + {"data": [{"image": FILE_PATH, WSIPatchKeys.LEVEL: 8, WSIPatchKeys.SIZE: (2, 2)}], "mask_level": 8}, + { + "num_patches": 4256, + "wsi_size": [32914, 46000], + "mask_level": 8, + "patch_level": 8, + "mask_size": (128, 179), + "patch_size": (2, 2), + }, +] + + +@skipUnless(has_cucim or has_osl or has_tiff, "Requires cucim, openslide, or tifffile!") +def setUpModule(): + hash_type = testing_data_config("images", FILE_KEY, "hash_type") + hash_val = testing_data_config("images", FILE_KEY, "hash_val") + download_url_or_skip_test(FILE_URL, FILE_PATH, hash_type=hash_type, hash_val=hash_val) + + +class MaskedPatchWSIDatasetTests: + class Tests(unittest.TestCase): + backend = None + + @parameterized.expand([TEST_CASE_0]) + def test_gen_patches(self, input_parameters, expected): + dataset = MaskedPatchWSIDataset(reader=self.backend, **input_parameters) + self.assertEqual(len(dataset), expected["num_patches"]) + for i, sample in enumerate(dataset): + self.assertEqual(sample["metadata"][WSIPatchKeys.LEVEL], expected["patch_level"]) + assert_array_equal(sample["metadata"][WSIPatchKeys.SIZE], expected["patch_size"]) + assert_array_equal(sample["image"].shape[1:], expected["patch_size"]) + self.assertTrue(sample["metadata"][WSIPatchKeys.LOCATION][0] >= 0) + self.assertTrue(sample["metadata"][WSIPatchKeys.LOCATION][0] < expected["wsi_size"][0]) + self.assertTrue(sample["metadata"][WSIPatchKeys.LOCATION][1] >= 0) + self.assertTrue(sample["metadata"][WSIPatchKeys.LOCATION][1] < expected["wsi_size"][1]) + if i > 10: + break + + +@skipUnless(has_cucim, "Requires cucim") +class TestSlidingPatchWSIDatasetCuCIM(MaskedPatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "cucim" + + +@skipUnless(has_osl, "Requires openslide") +class TestSlidingPatchWSIDatasetOpenSlide(MaskedPatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "openslide" + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_matshow3d.py b/tests/test_matshow3d.py index 83984d1556..18ed16dd2c 100644 --- a/tests/test_matshow3d.py +++ b/tests/test_matshow3d.py @@ -42,6 +42,9 @@ def test_3d(self): comp = compare_images(f"{testing_dir}/matshow3d_test.png", tempimg, 5e-2) self.assertIsNone(comp, f"value of comp={comp}") # None indicates test passed + _, axes = pyplot.subplots() + matshow3d(ims[keys], fig=axes, figsize=(2, 2), frames_per_row=5, every_n=2, frame_dim=-1, show=False) + def test_samples(self): testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") keys = "image" diff --git a/tests/test_mednistdataset.py b/tests/test_mednistdataset.py index e7cc1a60ff..87a0d6a2ec 100644 --- a/tests/test_mednistdataset.py +++ b/tests/test_mednistdataset.py @@ -15,8 +15,8 @@ from pathlib import Path from monai.apps import MedNISTDataset -from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd, ToTensord -from monai.utils.enums import PostFix +from monai.data import MetaTensor +from monai.transforms import AddChanneld, Compose, LoadImaged, ScaleIntensityd from tests.utils import skip_if_downloading_fails, skip_if_quick MEDNIST_FULL_DATASET_LENGTH = 58954 @@ -26,20 +26,13 @@ class TestMedNISTDataset(unittest.TestCase): @skip_if_quick def test_values(self): testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") - transform = Compose( - [ - LoadImaged(keys="image"), - AddChanneld(keys="image"), - ScaleIntensityd(keys="image"), - ToTensord(keys=["image", "label"]), - ] - ) + transform = Compose([LoadImaged(keys="image"), AddChanneld(keys="image"), ScaleIntensityd(keys="image")]) def _test_dataset(dataset): self.assertEqual(len(dataset), int(MEDNIST_FULL_DATASET_LENGTH * dataset.test_frac)) self.assertTrue("image" in dataset[0]) self.assertTrue("label" in dataset[0]) - self.assertTrue(PostFix.meta("image") in dataset[0]) + self.assertTrue(isinstance(dataset[0]["image"], MetaTensor)) self.assertTupleEqual(dataset[0]["image"].shape, (1, 64, 64)) with skip_if_downloading_fails(): @@ -59,7 +52,7 @@ def _test_dataset(dataset): data = MedNISTDataset(root_dir=testing_dir, transform=transform, section="test", download=False, seed=42) _test_dataset(data) self.assertEqual(data[0]["class_name"], "AbdomenCT") - self.assertEqual(data[0]["label"].cpu().item(), 0) + self.assertEqual(data[0]["label"], 0) shutil.rmtree(os.path.join(testing_dir, "MedNIST")) try: MedNISTDataset(root_dir=testing_dir, transform=transform, section="test", download=False) diff --git a/tests/test_meta_affine.py b/tests/test_meta_affine.py new file mode 100644 index 0000000000..269db33ef4 --- /dev/null +++ b/tests/test_meta_affine.py @@ -0,0 +1,179 @@ +# 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. + +import os +import unittest +from copy import deepcopy + +import numpy as np +from parameterized import parameterized + +from monai.data.image_writer import ITKWriter +from monai.transforms import ( + Compose, + EnsureChannelFirst, + EnsureChannelFirstd, + LoadImage, + LoadImaged, + MapTransform, + Orientation, + Orientationd, + Randomizable, + Spacing, + Spacingd, + Transform, +) +from monai.utils import convert_data_type, optional_import +from tests.utils import assert_allclose, download_url_or_skip_test, testing_data_config + +itk, has_itk = optional_import("itk") +TINY_DIFF = 1e-4 + +keys = ("img1", "img2") +key, key_1 = "ref_avg152T1_LR", "ref_avg152T1_RL" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", f"{key}.nii.gz") +FILE_PATH_1 = os.path.join(os.path.dirname(__file__), "testing_data", f"{key_1}.nii.gz") + +TEST_CASES_ARRAY = [ + [Compose([Spacing(pixdim=(1.0, 1.1, 1.2)), Orientation(axcodes="RAS")]), {}, TINY_DIFF], + [Compose([Orientation(axcodes="RAS"), Spacing(pixdim=(1.0, 1.1, 1.2))]), {}, TINY_DIFF], + ["CropForeground", {"k_divisible": 3}, TINY_DIFF], + ["BorderPad", {"spatial_border": (2, 3, 4)}, TINY_DIFF], + ["CenterScaleCrop", {"roi_scale": (0.6, 0.7, 0.8)}, TINY_DIFF], + ["CenterSpatialCrop", {"roi_size": (30, 200, 52)}, TINY_DIFF], + ["DivisiblePad", {"k": 16}, TINY_DIFF], + ["RandScaleCrop", {"roi_scale": (0.3, 0.4, 0.5)}, TINY_DIFF], + ["RandSpatialCrop", {"roi_size": (31, 32, 33)}, TINY_DIFF], + ["ResizeWithPadOrCrop", {"spatial_size": (50, 80, 200)}, TINY_DIFF], + ["Spacing", {"pixdim": (1.0, 1.1, 1.2)}, TINY_DIFF], + ["Orientation", {"axcodes": "RAS"}, TINY_DIFF], + ["Flip", {"spatial_axis": (0, 1)}, TINY_DIFF], + ["Resize", {"spatial_size": (100, 201, 1)}, 30.0], + ["Rotate", {"angle": (np.pi / 3, np.pi / 2, np.pi / 4), "mode": "bilinear"}, 20.0], + ["Zoom", {"zoom": (0.8, 0.91, 1.2)}, 20.0], + ["Rotate90", {"k": 3}, TINY_DIFF], + ["RandRotate90", {"prob": 1.0, "max_k": 3}, TINY_DIFF], + ["RandRotate", {"prob": 1.0, "range_x": np.pi / 3}, 20.0], + ["RandFlip", {"prob": 1.0}, TINY_DIFF], + ["RandAxisFlip", {"prob": 1.0}, TINY_DIFF], + ["RandZoom", {"prob": 1.0, "mode": "trilinear"}, TINY_DIFF], + [ + "RandAffine", + { + "prob": 1.0, + "rotate_range": (np.pi / 4, np.pi / 3, np.pi / 2), + "translate_range": (3, 4, 5), + "scale_range": (-0.1, 0.2), + "spatial_size": (30, 40, 50), + "cache_grid": True, + "mode": "bilinear", + }, + 20.0, + ], + [ + "Affine", + { + "rotate_params": (np.pi / 4, 0.0, 0.0), + "translate_params": (3, 4, 5), + "spatial_size": (30, 40, 50), + "mode": "bilinear", + "image_only": True, + }, + 20.0, + ], +] + +TEST_CASES_DICT = [ + [Compose([Spacingd(keys, pixdim=(1.0, 1.1, 1.2)), Orientationd(keys, axcodes="LAS")]), {}, TINY_DIFF], + [Compose([Orientationd(keys, axcodes="LAS"), Spacingd(keys, pixdim=(1.0, 1.1, 1.2))]), {}, TINY_DIFF], + ["CropForegroundd", {"k_divisible": 3, "source_key": "img1"}, TINY_DIFF], +] +for c in TEST_CASES_ARRAY[3:-1]: # exclude CropForegroundd and Affined + TEST_CASES_DICT.append(deepcopy(c)) + TEST_CASES_DICT[-1][0] = TEST_CASES_DICT[-1][0] + "d" # type: ignore + + +def _create_itk_obj(array, affine): + itk_img = deepcopy(array) + itk_img = convert_data_type(itk_img, np.ndarray)[0] + itk_obj = ITKWriter.create_backend_obj(itk_img, channel_dim=None, affine=affine, affine_lps=True) + return itk_obj + + +def _resample_to_affine(itk_obj, ref_obj): + """linear resample""" + dim = itk_obj.GetImageDimension() + transform = itk.IdentityTransform[itk.D, dim].New() + interpolator = itk.LinearInterpolateImageFunction[type(itk_obj), itk.D].New() + resampled = itk.resample_image_filter( + Input=itk_obj, interpolator=interpolator, transform=transform, UseReferenceImage=True, ReferenceImage=ref_obj + ) + return resampled + + +@unittest.skipUnless(has_itk, "Requires itk package.") +class TestAffineConsistencyITK(unittest.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + for k, n in ((key, FILE_PATH), (key_1, FILE_PATH_1)): + config = testing_data_config("images", f"{k}") + download_url_or_skip_test(filepath=n, **config) + + def run_transform(self, img, xform_cls, args_dict): + if isinstance(xform_cls, Transform): + xform = xform_cls + output = xform(img, **args_dict) + else: + if isinstance(xform_cls, str): + xform_cls, _ = optional_import("monai.transforms", name=xform_cls) + if issubclass(xform_cls, MapTransform): + args_dict.update({"keys": keys}) + xform = xform_cls(**args_dict) + if isinstance(xform, Randomizable): + xform.set_random_state(5) + output = xform(img) + return output + + @parameterized.expand(TEST_CASES_ARRAY) + def test_linear_consistent(self, xform_cls, input_dict, atol): + """xform cls testing itk consistency""" + img = LoadImage(image_only=True, simple_keys=True)(FILE_PATH) + img = EnsureChannelFirst()(img) + ref_1 = _create_itk_obj(img[0], img.affine) + output = self.run_transform(img, xform_cls, input_dict) + ref_2 = _create_itk_obj(output[0], output.affine) + assert_allclose(output.pixdim, np.asarray(ref_2.GetSpacing()), type_test=False) + expected = _resample_to_affine(ref_1, ref_2) + # compare ref_2 and expected results from itk + diff = np.abs(itk.GetArrayFromImage(ref_2) - itk.GetArrayFromImage(expected)) + avg_diff = np.mean(diff) + + self.assertTrue(avg_diff < atol, f"{xform_cls} avg_diff: {avg_diff}, tol: {atol}") + + @parameterized.expand(TEST_CASES_DICT) + def test_linear_consistent_dict(self, xform_cls, input_dict, atol): + """xform cls testing itk consistency""" + img = LoadImaged(keys, image_only=True, simple_keys=True)({keys[0]: FILE_PATH, keys[1]: FILE_PATH_1}) + img = EnsureChannelFirstd(keys)(img) + ref_1 = {k: _create_itk_obj(img[k][0], img[k].affine) for k in keys} + output = self.run_transform(img, xform_cls, input_dict) + ref_2 = {k: _create_itk_obj(output[k][0], output[k].affine) for k in keys} + expected = {k: _resample_to_affine(ref_1[k], ref_2[k]) for k in keys} + # compare ref_2 and expected results from itk + diff = {k: np.abs(itk.GetArrayFromImage(ref_2[k]) - itk.GetArrayFromImage(expected[k])) for k in keys} + avg_diff = {k: np.mean(diff[k]) for k in keys} + for k in keys: + self.assertTrue(avg_diff[k] < atol, f"{xform_cls} avg_diff: {avg_diff}, tol: {atol}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py new file mode 100644 index 0000000000..c1b8bc5046 --- /dev/null +++ b/tests/test_meta_tensor.py @@ -0,0 +1,524 @@ +# 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. + +import io +import os +import random +import string +import tempfile +import unittest +import warnings +from copy import deepcopy +from multiprocessing.reduction import ForkingPickler +from typing import Optional, Union + +import numpy as np +import torch +import torch.multiprocessing +from parameterized import parameterized + +from monai.data import DataLoader, Dataset +from monai.data.meta_obj import get_track_meta, set_track_meta +from monai.data.meta_tensor import MetaTensor +from monai.data.utils import decollate_batch, list_data_collate +from monai.transforms import BorderPadd, Compose, DivisiblePadd, FromMetaTensord, ToMetaTensord +from monai.utils.enums import PostFix +from monai.utils.module import pytorch_after +from tests.utils import TEST_DEVICES, SkipIfBeforePyTorchVersion, assert_allclose, skip_if_no_cuda + +DTYPES = [[torch.float32], [torch.float64], [torch.float16], [torch.int64], [torch.int32], [None]] +TESTS = [] +for _device in TEST_DEVICES: + for _dtype in DTYPES: + TESTS.append((*_device, *_dtype)) # type: ignore + + +def rand_string(min_len=5, max_len=10): + str_size = random.randint(min_len, max_len) + chars = string.ascii_letters + string.punctuation + return "".join(random.choice(chars) for _ in range(str_size)) + + +class TestMetaTensor(unittest.TestCase): + @staticmethod + def get_im(shape=None, dtype=None, device=None): + if shape is None: + shape = (1, 10, 8) + affine = torch.randint(0, 10, (4, 4)) + meta = {"fname": rand_string()} + t = torch.rand(shape) + if dtype is not None: + t = t.to(dtype) + if device is not None: + t = t.to(device) + m = MetaTensor(t.clone(), affine, meta) + return m, t + + def check_ids(self, a, b, should_match): + comp = self.assertEqual if should_match else self.assertNotEqual + comp(id(a), id(b)) + + def check_meta(self, a: MetaTensor, b: MetaTensor) -> None: + self.assertEqual(a.is_batch, b.is_batch) + meta_a, meta_b = a.meta, b.meta + # need to split affine from rest of metadata + aff_a = meta_a.get("affine", None) + aff_b = meta_b.get("affine", None) + assert_allclose(aff_a, aff_b) + meta_a = {k: v for k, v in meta_a.items() if k != "affine"} + meta_b = {k: v for k, v in meta_b.items() if k != "affine"} + self.assertEqual(meta_a, meta_b) + + def check( + self, + out: torch.Tensor, + orig: torch.Tensor, + *, + shape: bool = True, + vals: bool = True, + ids: bool = True, + device: Optional[Union[str, torch.device]] = None, + meta: bool = True, + check_ids: bool = True, + **kwargs, + ): + if device is None: + device = orig.device + + # check the image + self.assertIsInstance(out, type(orig)) + if shape: + assert_allclose(torch.as_tensor(out.shape), torch.as_tensor(orig.shape)) + if vals: + assert_allclose(out, orig, **kwargs) + if check_ids: + self.check_ids(out, orig, ids) + self.assertTrue(str(device) in str(out.device)) + + # check meta and affine are equal and affine is on correct device + if isinstance(orig, MetaTensor) and isinstance(out, MetaTensor) and meta: + self.check_meta(orig, out) + self.assertTrue(str(device) in str(out.affine.device)) + if check_ids: + self.check_ids(out.affine, orig.affine, ids) + self.check_ids(out.meta, orig.meta, ids) + + @parameterized.expand(TESTS) + def test_as_tensor(self, device, dtype): + m, t = self.get_im(device=device, dtype=dtype) + t2 = m.as_tensor() + self.assertIsInstance(t2, torch.Tensor) + self.assertNotIsInstance(t2, MetaTensor) + self.assertIsInstance(m, MetaTensor) + self.check(t, t2, ids=False) + + def test_as_dict(self): + m, _ = self.get_im() + m_dict = m.as_dict("im") + im, meta = m_dict["im"], m_dict[PostFix.meta("im")] + affine = meta.pop("affine") + m2 = MetaTensor(im, affine, meta) + self.check(m2, m, check_ids=False) + + @parameterized.expand(TESTS) + def test_constructor(self, device, dtype): + m, t = self.get_im(device=device, dtype=dtype) + # construct from pre-existing + m1 = MetaTensor(m.clone()) + self.check(m, m1, ids=False, meta=False) + # meta already has affine + m2 = MetaTensor(t.clone(), meta=m.meta) + self.check(m, m2, ids=False, meta=False) + # meta dosen't have affine + affine = m.meta.pop("affine") + m3 = MetaTensor(t.clone(), affine=affine, meta=m.meta) + self.check(m, m3, ids=False, meta=False) + + @parameterized.expand(TESTS) + @skip_if_no_cuda + def test_to_cuda(self, device, dtype): + """Test `to`, `cpu` and `cuda`. For `to`, check args and kwargs.""" + orig, _ = self.get_im(device=device, dtype=dtype) + m = orig.clone() + m = m.to("cuda") + self.check(m, orig, ids=False, device="cuda") + m = m.cpu() + self.check(m, orig, ids=False, device="cpu") + m = m.cuda() + self.check(m, orig, ids=False, device="cuda") + m = m.to("cpu") + self.check(m, orig, ids=False, device="cpu") + m = m.to(device="cuda") + self.check(m, orig, ids=False, device="cuda") + m = m.to(device="cpu") + self.check(m, orig, ids=False, device="cpu") + + @skip_if_no_cuda + def test_affine_device(self): + m, _ = self.get_im() # device="cuda") + m.affine = torch.eye(4) + self.assertEqual(m.device, m.affine.device) + + @parameterized.expand(TESTS) + def test_copy(self, device, dtype): + m, _ = self.get_im(device=device, dtype=dtype) + # shallow copy + a = m + self.check(a, m, ids=True) + # deepcopy + a = deepcopy(m) + self.check(a, m, ids=False) + # clone + a = m.clone() + self.check(a, m, ids=False) + + @parameterized.expand(TESTS) + def test_add(self, device, dtype): + m1, t1 = self.get_im(device=device, dtype=dtype) + m2, t2 = self.get_im(device=device, dtype=dtype) + self.check(m1 + m2, t1 + t2, ids=False) + self.check(torch.add(m1, m2), t1 + t2, ids=False) + self.check(torch.add(input=m1, other=m2), t1 + t2, ids=False) + self.check(torch.add(m1, other=m2), t1 + t2, ids=False) + m3 = deepcopy(m2) + t3 = deepcopy(t2) + m3 += 3 + t3 += 3 + self.check(m3, t3, ids=False) + # check torch.Tensor+MetaTensor and MetaTensor+torch.Tensor + self.check(torch.add(m1, t2), t1 + t2, ids=False) + self.check(torch.add(t2, m1), t1 + t2, ids=False) + + @parameterized.expand(TEST_DEVICES) + def test_conv(self, device): + im, _ = self.get_im((1, 3, 10, 8, 12), device=device) + conv = torch.nn.Conv3d(im.shape[1], 5, 3) + conv.to(device) + out = conv(im) + self.check(out, im, shape=False, vals=False, ids=False) + + @parameterized.expand(TESTS) + def test_stack(self, device, dtype): + numel = 3 + ims = [self.get_im(device=device, dtype=dtype)[0] for _ in range(numel)] + stacked = torch.stack(ims) + self.assertIsInstance(stacked, MetaTensor) + orig_affine = ims[0].meta.pop("affine") + stacked_affine = stacked.meta.pop("affine") + assert_allclose(orig_affine, stacked_affine) + self.assertEqual(stacked.meta, ims[0].meta) + + def test_get_set_meta_fns(self): + set_track_meta(False) + self.assertEqual(get_track_meta(), False) + set_track_meta(True) + self.assertEqual(get_track_meta(), True) + + @parameterized.expand(TEST_DEVICES) + def test_torchscript(self, device): + shape = (1, 3, 10, 8) + im, _ = self.get_im(shape, device=device) + conv = torch.nn.Conv2d(im.shape[1], 5, 3) + conv.to(device) + im_conv = conv(im) + traced_fn = torch.jit.trace(conv, im.as_tensor()) + # save it, load it, use it + with tempfile.TemporaryDirectory() as tmp_dir: + fname = os.path.join(tmp_dir, "im.pt") + torch.jit.save(traced_fn, f=fname) + traced_fn = torch.jit.load(fname) + out = traced_fn(im) + self.assertIsInstance(out, torch.Tensor) + if not isinstance(out, MetaTensor) and not pytorch_after(1, 9, 1): + warnings.warn( + "When calling `nn.Module(MetaTensor) on a module traced with " + "`torch.jit.trace`, your version of pytorch returns a " + "`torch.Tensor` instead of a `MetaTensor`. Consider upgrading " + "your pytorch version if this is important to you." + ) + im_conv = im_conv.as_tensor() + self.check(out, im_conv, ids=False) + + def test_pickling(self): + m, _ = self.get_im() + with tempfile.TemporaryDirectory() as tmp_dir: + fname = os.path.join(tmp_dir, "im.pt") + torch.save(m, fname) + m2 = torch.load(fname) + if not isinstance(m2, MetaTensor) and not pytorch_after(1, 8, 1): + warnings.warn("Old version of pytorch. pickling converts `MetaTensor` to `torch.Tensor`.") + m = m.as_tensor() + self.check(m2, m, ids=False) + + @skip_if_no_cuda + def test_amp(self): + shape = (1, 3, 10, 8) + device = "cuda" + im, _ = self.get_im(shape, device=device) + conv = torch.nn.Conv2d(im.shape[1], 5, 3) + conv.to(device) + im_conv = conv(im) + with torch.cuda.amp.autocast(): + im_conv2 = conv(im) + self.check(im_conv2, im_conv, ids=False, rtol=1e-2, atol=1e-2) + + def test_out(self): + """Test when `out` is given as an argument.""" + m1, _ = self.get_im() + m2, _ = self.get_im() + m3, _ = self.get_im() + torch.add(m2, m3, out=m1) + m1_add = m2 + m3 + + assert_allclose(m1, m1_add) + # self.check_meta(m1, m2) # meta is from first input tensor + + @parameterized.expand(TESTS) + def test_collate(self, device, dtype): + numel = 3 + ims = [self.get_im(device=device, dtype=dtype)[0] for _ in range(numel)] + ims = [MetaTensor(im, applied_operations=[f"t{i}"]) for i, im in enumerate(ims)] + collated = list_data_collate(ims) + # tensor + self.assertIsInstance(collated, MetaTensor) + expected_shape = (numel,) + tuple(ims[0].shape) + self.assertTupleEqual(tuple(collated.shape), expected_shape) + for i, im in enumerate(ims): + self.check(im, ims[i], ids=True) + # affine + self.assertIsInstance(collated.affine, torch.Tensor) + expected_shape = (numel,) + tuple(ims[0].affine.shape) + self.assertTupleEqual(tuple(collated.affine.shape), expected_shape) + self.assertEqual(len(collated.applied_operations), numel) + + @parameterized.expand(TESTS) + def test_dataset(self, device, dtype): + ims = [self.get_im(device=device, dtype=dtype)[0] for _ in range(4)] + ds = Dataset(ims) + for i, im in enumerate(ds): + self.check(im, ims[i], ids=True) + + @parameterized.expand(DTYPES) + @SkipIfBeforePyTorchVersion((1, 8)) + def test_dataloader(self, dtype): + batch_size = 5 + ims = [self.get_im(dtype=dtype)[0] for _ in range(batch_size * 2)] + ims = [MetaTensor(im, applied_operations=[f"t{i}"]) for i, im in enumerate(ims)] + ds = Dataset(ims) + im_shape = tuple(ims[0].shape) + affine_shape = tuple(ims[0].affine.shape) + expected_im_shape = (batch_size,) + im_shape + expected_affine_shape = (batch_size,) + affine_shape + dl = DataLoader(ds, num_workers=batch_size, batch_size=batch_size) + for batch in dl: + self.assertIsInstance(batch, MetaTensor) + self.assertTupleEqual(tuple(batch.shape), expected_im_shape) + self.assertTupleEqual(tuple(batch.affine.shape), expected_affine_shape) + self.assertEqual(len(batch.applied_operations), batch_size) + + @SkipIfBeforePyTorchVersion((1, 9)) + def test_indexing(self): + """ + Check the metadata is returned in the expected format depending on whether + the input `MetaTensor` is a batch of data or not. + """ + ims = [self.get_im()[0] for _ in range(5)] + data = list_data_collate(ims) + + # check that when using non-batch data, metadata is copied wholly when indexing + # or iterating across data. + im = ims[0] + self.check_meta(im[0], im) + self.check_meta(next(iter(im)), im) + + self.assertEqual(im[None].shape, (1, 1, 10, 8)) + self.assertEqual(data[None].shape, (1, 5, 1, 10, 8)) + + # index + d = data[0] + self.check(d, ims[0], ids=False) + + # iter + d = next(iter(data)) + self.check(d, ims[0], ids=False) + + # complex indexing + + # `is_batch==True`, should have subset of image and metadata. + d = data[1:3] + self.check(d, list_data_collate(ims[1:3]), ids=False) + + # is_batch==True, should have subset of image and same metadata as `[1:3]`. + d = data[1:3, 0] + self.check(d, list_data_collate([i[0] for i in ims[1:3]]), ids=False) + + # `is_batch==False`, should have first metadata and subset of first image. + d = data[0, 0] + self.check(d, ims[0][0], ids=False) + + # `is_batch==True`, should have all metadata and subset of all images. + d = data[:, 0] + self.check(d, list_data_collate([i[0] for i in ims]), ids=False) + + # `is_batch==True`, should have all metadata and subset of all images. + d = data[..., -1] + self.check(d, list_data_collate([i[..., -1] for i in ims]), ids=False) + + # `is_batch==False`, tuple split along batch dim. Should have individual + # metadata. + d = data.unbind(0) + self.assertIsInstance(d, tuple) + self.assertEqual(len(d), len(ims)) + for _d, _im in zip(d, ims): + self.check(_d, _im, ids=False) + + # `is_batch==False`, tuple split along batch dim. Should have individual + # metadata. + d = data.unbind(dim=0) + self.assertIsInstance(d, tuple) + self.assertEqual(len(d), len(ims)) + for _d, _im in zip(d, ims): + self.check(_d, _im, ids=False) + + # `is_batch==True`, tuple split along non-batch dim. Should have all metadata. + d = data.unbind(-1) + self.assertIsInstance(d, tuple) + self.assertEqual(len(d), ims[0].shape[-1]) + for _d in d: + self.check_meta(_d, data) + + # `is_batch==True`, tuple split along non-batch dim. Should have all metadata. + d = data.unbind(dim=-1) + self.assertIsInstance(d, tuple) + self.assertEqual(len(d), ims[0].shape[-1]) + for _d in d: + self.check_meta(_d, data) + + @parameterized.expand(DTYPES) + @SkipIfBeforePyTorchVersion((1, 8)) + def test_decollate(self, dtype): + batch_size = 3 + ims = [self.get_im(dtype=dtype)[0] for _ in range(batch_size * 2)] + ds = Dataset(ims) + dl = DataLoader(ds, num_workers=batch_size, batch_size=batch_size) + batch = next(iter(dl)) + decollated = decollate_batch(batch) + self.assertIsInstance(decollated, list) + self.assertEqual(len(decollated), batch_size) + for elem, im in zip(decollated, ims): + self.assertIsInstance(elem, MetaTensor) + self.check(elem, im, ids=False) + + def test_str(self): + t = MetaTensor([1.0], affine=torch.tensor(1), meta={"fname": "filename"}) + s1 = str(t) + s2 = t.__repr__() + expected_out = ( + "tensor([1.])\n" + + "MetaData\n" + + "\tfname: filename\n" + + "\taffine: 1\n" + + "\n" + + "Applied operations\n" + + "[]\n" + + "Is batch?: False" + ) + for s in (s1, s2): + self.assertEqual(s, expected_out) + + def test_astype(self): + t = MetaTensor([1.0], affine=torch.tensor(1), meta={"fname": "filename"}) + for np_types in ("float32", "np.float32", "numpy.float32", np.float32, float, "int", np.compat.long, 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) + + def test_transforms(self): + key = "im" + _, im = self.get_im() + tr = Compose([ToMetaTensord(key), BorderPadd(key, 1), DivisiblePadd(key, 16), FromMetaTensord(key)]) + num_tr = len(tr.transforms) + data = {key: im, PostFix.meta(key): {"affine": torch.eye(4)}} + + # apply one at a time + for i, _tr in enumerate(tr.transforms): + data = _tr(data) + is_meta = isinstance(_tr, (ToMetaTensord, BorderPadd, DivisiblePadd)) + if is_meta: + self.assertEqual(len(data), 1) # im + self.assertIsInstance(data[key], MetaTensor) + n_applied = len(data[key].applied_operations) + else: + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + + self.assertEqual(n_applied, i + 1) + + # inverse one at a time + for i, _tr in enumerate(tr.transforms[::-1]): + data = _tr.inverse(data) + is_meta = isinstance(_tr, (FromMetaTensord, BorderPadd, DivisiblePadd)) + if is_meta: + self.assertEqual(len(data), 1) # im + self.assertIsInstance(data[key], MetaTensor) + n_applied = len(data[key].applied_operations) + else: + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + + self.assertEqual(n_applied, num_tr - i - 1) + + # apply all in one go + data = tr({key: im, PostFix.meta(key): {"affine": torch.eye(4)}}) + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + self.assertEqual(n_applied, num_tr) + + # inverse all in one go + data = tr.inverse(data) + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + self.assertEqual(n_applied, 0) + + def test_construct_with_pre_applied_transforms(self): + key = "im" + _, im = self.get_im() + tr = Compose([BorderPadd(key, 1), DivisiblePadd(key, 16)]) + data = tr({key: im}) + m = MetaTensor(im, applied_operations=data["im"].applied_operations) + self.assertEqual(len(m.applied_operations), len(tr.transforms)) + + @parameterized.expand(TESTS) + def test_multiprocessing(self, device=None, dtype=None): + """multiprocessing sharing with 'device' and 'dtype'""" + buf = io.BytesIO() + t = MetaTensor([0.0, 0.0], device=device, dtype=dtype) + if t.is_cuda: + with self.assertRaises(NotImplementedError): + ForkingPickler(buf).dump(t) + return + ForkingPickler(buf).dump(t) + obj = ForkingPickler.loads(buf.getvalue()) + self.assertIsInstance(obj, MetaTensor) + assert_allclose(obj.as_tensor(), t) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_metatensor_integration.py b/tests/test_metatensor_integration.py new file mode 100644 index 0000000000..d6908815ee --- /dev/null +++ b/tests/test_metatensor_integration.py @@ -0,0 +1,101 @@ +# 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. + +import os +import tempfile +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.bundle import ConfigParser +from monai.data import CacheDataset, DataLoader, MetaTensor, decollate_batch +from monai.data.utils import TraceKeys +from monai.transforms import InvertD, SaveImageD +from monai.utils import optional_import, set_determinism +from tests.utils import assert_allclose, download_url_or_skip_test, testing_data_config + +nib, has_nib = optional_import("nibabel") +TINY_DIFF = 0.1 + +keys = ("img", "seg") +key, key_1 = "MNI152_T1_2mm", "MNI152_T1_2mm_strucseg" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", f"{key}.nii.gz") +FILE_PATH_1 = os.path.join(os.path.dirname(__file__), "testing_data", f"{key_1}.nii.gz") +TEST_CASES = os.path.join(os.path.dirname(__file__), "testing_data", "transform_metatensor_cases.yaml") + + +@unittest.skipUnless(has_nib, "Requires nibabel package.") +class TestMetaTensorIntegration(unittest.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + for k, n in ((key, FILE_PATH), (key_1, FILE_PATH_1)): + config = testing_data_config("images", f"{k}") + download_url_or_skip_test(filepath=n, **config) + cls.files = [{keys[0]: x, keys[1]: y} for (x, y) in [[FILE_PATH, FILE_PATH_1]] * 4] + + @classmethod + def tearDownClass(cls): + super().tearDownClass() + set_determinism(None) + + @parameterized.expand(["TEST_CASE_1", "TEST_CASE_2", "TEST_CASE_3"]) + def test_transforms(self, case_id): + set_determinism(2022) + config = ConfigParser() + config.read_config(TEST_CASES) + config["input_keys"] = keys + test_case = config.get_parsed_content(id=case_id, instantiate=True) # transform instance + + dataset = CacheDataset(self.files, transform=test_case) + loader = DataLoader(dataset, batch_size=3, shuffle=True) + for x in loader: + self.assertIsInstance(x[keys[0]], MetaTensor) + self.assertIsInstance(x[keys[1]], MetaTensor) + out = decollate_batch(x) # decollate every batch should work + + # test forward patches + loaded = out[0] + self.assertEqual(len(loaded), len(keys)) + img, seg = loaded[keys[0]], loaded[keys[1]] + expected = config.get_parsed_content(id=f"{case_id}_answer", instantiate=True) # expected results + self.assertEqual(expected["load_shape"], list(x[keys[0]].shape)) + assert_allclose(expected["affine"], img.affine, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + assert_allclose(expected["affine"], seg.affine, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + test_cls = [type(x).__name__ for x in test_case.transforms] + tracked_cls = [x[TraceKeys.CLASS_NAME] for x in img.applied_operations] + self.assertTrue(len(tracked_cls) <= len(test_cls)) # tracked items should be no more than the compose items. + with tempfile.TemporaryDirectory() as tempdir: # test writer + SaveImageD(keys, resample=False, output_dir=tempdir, output_postfix=case_id)(loaded) + + # test inverse + inv = InvertD(keys, orig_keys=keys, transform=test_case, nearest_interp=True) + out = inv(loaded) + img, seg = out[keys[0]], out[keys[1]] + assert_allclose(expected["inv_affine"], img.affine, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + assert_allclose(expected["inv_affine"], seg.affine, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + self.assertFalse(img.applied_operations) + self.assertFalse(seg.applied_operations) + assert_allclose(expected["inv_shape"], img.shape, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + assert_allclose(expected["inv_shape"], seg.shape, type_test=False, atol=TINY_DIFF, rtol=TINY_DIFF) + with tempfile.TemporaryDirectory() as tempdir: # test writer + SaveImageD(keys, resample=False, output_dir=tempdir, output_postfix=case_id)(out) + seg_file = os.path.join(tempdir, key_1, f"{key_1}_{case_id}.nii.gz") + segout = nib.load(seg_file).get_fdata() + segin = nib.load(FILE_PATH_1).get_fdata() + ndiff = np.sum(np.abs(segout - segin) > 0) + total = np.prod(segout.shape) + self.assertTrue(ndiff / total < 0.4, f"{ndiff / total}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mlp.py b/tests/test_mlp.py index 6fec5b6854..737762cfb1 100644 --- a/tests/test_mlp.py +++ b/tests/test_mlp.py @@ -21,7 +21,7 @@ TEST_CASE_MLP = [] for dropout_rate in np.linspace(0, 1, 4): for hidden_size in [128, 256, 512, 768]: - for mlp_dim in [512, 1028, 2048, 3072]: + for mlp_dim in [0, 1028, 2048, 3072]: test_case = [ {"hidden_size": hidden_size, "mlp_dim": mlp_dim, "dropout_rate": dropout_rate}, diff --git a/tests/test_module_list.py b/tests/test_module_list.py index 83c6979f30..d0b5aaf26b 100644 --- a/tests/test_module_list.py +++ b/tests/test_module_list.py @@ -38,8 +38,9 @@ def test_public_api(self): def test_transform_api(self): """monai subclasses of MapTransforms must have alias names ending with 'd', 'D', 'Dict'""" to_exclude = {"MapTransform"} # except for these transforms - to_exclude_docs = {"Decollate", "Ensemble", "Invert", "SaveClassification", "RandTorchVision"} + to_exclude_docs = {"Decollate", "Ensemble", "Invert", "SaveClassification", "RandTorchVision", "RandCrop"} to_exclude_docs.update({"DeleteItems", "SelectItems", "CopyItems", "ConcatItems"}) + to_exclude_docs.update({"ToMetaTensor", "FromMetaTensor"}) xforms = { name: obj for name, obj in monai.transforms.__dict__.items() diff --git a/tests/test_multi_scale.py b/tests/test_multi_scale.py index 963824f25e..f348c09512 100644 --- a/tests/test_multi_scale.py +++ b/tests/test_multi_scale.py @@ -16,7 +16,7 @@ from monai.losses import DiceLoss from monai.losses.multi_scale import MultiScaleLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save dice_loss = DiceLoss(include_background=True, sigmoid=True, smooth_nr=1e-5, smooth_dr=1e-5) device = "cuda" if torch.cuda.is_available() else "cpu" @@ -67,7 +67,6 @@ def test_ill_opts(self): torch.ones((1, 1, 3), device=device), torch.ones((1, 1, 3), device=device) ) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): input_param, input_data, expected_val = TEST_CASES[0] loss = MultiScaleLoss(**input_param) diff --git a/tests/test_nifti_rw.py b/tests/test_nifti_rw.py index 2c0a8dc9a3..fae53394c3 100644 --- a/tests/test_nifti_rw.py +++ b/tests/test_nifti_rw.py @@ -30,14 +30,14 @@ [[-5.3, 0.0, 0.0, 102.01], [0.0, 0.52, 2.17, -7.50], [-0.0, 1.98, -0.26, -23.12], [0.0, 0.0, 0.0, 1.0]] ) ) - TESTS.append( - [ - TEST_IMAGE, - TEST_AFFINE, - dict(reader="NibabelReader", image_only=False, as_closest_canonical=True), - np.arange(24).reshape((2, 4, 3)), - ] - ) + # TESTS.append( + # [ + # TEST_IMAGE, + # TEST_AFFINE, + # dict(reader="NibabelReader", image_only=False, as_closest_canonical=True), + # np.arange(24).reshape((2, 4, 3)), + # ] + # ) TESTS.append( [ TEST_IMAGE, @@ -63,7 +63,7 @@ [ TEST_IMAGE, TEST_AFFINE, - dict(reader="NibabelReader", image_only=False, as_closest_canonical=False), + dict(reader="NibabelReader", image_only=True, as_closest_canonical=False), np.arange(24).reshape((2, 4, 3)), ] ) @@ -71,7 +71,7 @@ [ TEST_IMAGE, None, - dict(reader="NibabelReader", image_only=False, as_closest_canonical=False), + dict(reader="NibabelReader", image_only=True, as_closest_canonical=False), np.arange(24).reshape((2, 4, 3)), ] ) @@ -85,11 +85,12 @@ def test_orientation(self, array, affine, reader_param, expected): # read test cases loader = LoadImage(**reader_param) load_result = loader(test_image) - if isinstance(load_result, tuple): - data_array, header = load_result - else: - data_array = load_result + data_array = load_result.numpy() + if reader_param.get("image_only", False): header = None + else: + header = load_result.meta + header["affine"] = header["affine"].numpy() if os.path.exists(test_image): os.remove(test_image) @@ -114,9 +115,12 @@ def test_orientation(self, array, affine, reader_param, expected): def test_consistency(self): np.set_printoptions(suppress=True, precision=3) test_image = make_nifti_image(np.arange(64).reshape(1, 8, 8), np.diag([1.5, 1.5, 1.5, 1])) - data, header = LoadImage(reader="NibabelReader", as_closest_canonical=False)(test_image) - data, original_affine, new_affine = Spacing([0.8, 0.8, 0.8])(data[None], header["affine"], mode="nearest") - data, _, new_affine = Orientation("ILP")(data, new_affine) + data = LoadImage(image_only=True, reader="NibabelReader", as_closest_canonical=False)(test_image) + header = data.meta + data = Spacing([0.8, 0.8, 0.8])(data[None], header["affine"], mode="nearest") + original_affine = data.meta["original_affine"] + data = Orientation("ILP")(data) + new_affine = data.affine if os.path.exists(test_image): os.remove(test_image) writer_obj = NibabelWriter() diff --git a/tests/test_nifti_saver.py b/tests/test_nifti_saver.py index 6855a59041..54489904df 100644 --- a/tests/test_nifti_saver.py +++ b/tests/test_nifti_saver.py @@ -80,7 +80,7 @@ def test_saved_3d_no_resize_content(self): saver.save_batch(torch.randint(0, 255, (8, 8, 1, 2, 2)), meta_data) for i in range(8): filepath = os.path.join(tempdir, "testfile" + str(i), "testfile" + str(i) + "_seg.nii.gz") - img, _ = LoadImage("nibabelreader")(filepath) + img = LoadImage("nibabelreader", image_only=True)(filepath) self.assertEqual(img.shape, (1, 2, 2, 8)) def test_squeeze_end_dims(self): @@ -102,9 +102,8 @@ def test_squeeze_end_dims(self): # 2d image w channel saver.save(torch.randint(0, 255, (1, 2, 2)), meta_data) - im, meta = LoadImage()(os.path.join(tempdir, fname, fname + ".nii.gz")) + im = LoadImage(image_only=True)(os.path.join(tempdir, fname, fname + ".nii.gz")) self.assertTrue(im.ndim == 2 if squeeze_end_dims else 4) - self.assertTrue(meta["dim"][0] == im.ndim) if __name__ == "__main__": diff --git a/tests/test_normalize_intensity.py b/tests/test_normalize_intensity.py index 5bcee1263b..4d06a80c1d 100644 --- a/tests/test_normalize_intensity.py +++ b/tests/test_normalize_intensity.py @@ -86,19 +86,16 @@ def test_default(self, im_type): im = im_type(self.imt.copy()) normalizer = NormalizeIntensity() normalized = normalizer(im) - self.assertEqual(type(im), type(normalized)) - if isinstance(normalized, torch.Tensor): - self.assertEqual(im.device, normalized.device) self.assertTrue(normalized.dtype in (np.float32, torch.float32)) expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) - assert_allclose(normalized, expected, type_test=False, rtol=1e-3) + assert_allclose(normalized, expected, type_test="tensor", rtol=1e-3) @parameterized.expand(TESTS) def test_nonzero(self, in_type, input_param, input_data, expected_data): normalizer = NormalizeIntensity(**input_param) im = in_type(input_data) normalized = normalizer(im) - assert_allclose(normalized, in_type(expected_data)) + assert_allclose(normalized, in_type(expected_data), type_test="tensor") @parameterized.expand([[p] for p in TEST_NDARRAYS]) def test_channel_wise(self, im_type): @@ -106,7 +103,7 @@ def test_channel_wise(self, im_type): input_data = im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]])) expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]]) normalized = normalizer(input_data) - assert_allclose(normalized, im_type(expected)) + assert_allclose(normalized, im_type(expected), type_test="tensor") @parameterized.expand([[p] for p in TEST_NDARRAYS]) def test_value_errors(self, im_type): diff --git a/tests/test_normalize_intensityd.py b/tests/test_normalize_intensityd.py index 12a39b1b5b..a8167a1e93 100644 --- a/tests/test_normalize_intensityd.py +++ b/tests/test_normalize_intensityd.py @@ -12,7 +12,6 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import NormalizeIntensityd @@ -57,20 +56,14 @@ def test_image_normalize_intensityd(self, im_type): normalizer = NormalizeIntensityd(keys=[key]) normalized = normalizer({key: im})[key] expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) - self.assertEqual(type(im), type(normalized)) - if isinstance(normalized, torch.Tensor): - self.assertEqual(im.device, normalized.device) - assert_allclose(normalized, im_type(expected), rtol=1e-3) + assert_allclose(normalized, im_type(expected), rtol=1e-3, type_test="tensor") @parameterized.expand(TESTS) def test_nonzero(self, input_param, input_data, expected_data): key = "img" normalizer = NormalizeIntensityd(**input_param) normalized = normalizer(input_data)[key] - self.assertEqual(type(input_data[key]), type(normalized)) - if isinstance(normalized, torch.Tensor): - self.assertEqual(input_data[key].device, normalized.device) - assert_allclose(normalized, expected_data) + assert_allclose(normalized, expected_data, type_test="tensor") @parameterized.expand([[p] for p in TEST_NDARRAYS]) def test_channel_wise(self, im_type): @@ -78,11 +71,8 @@ def test_channel_wise(self, im_type): normalizer = NormalizeIntensityd(keys=key, nonzero=True, channel_wise=True) input_data = {key: im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]]))} normalized = normalizer(input_data)[key] - self.assertEqual(type(input_data[key]), type(normalized)) - if isinstance(normalized, torch.Tensor): - self.assertEqual(input_data[key].device, normalized.device) expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]]) - assert_allclose(normalized, im_type(expected)) + assert_allclose(normalized, im_type(expected), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_nrrd_reader.py b/tests/test_nrrd_reader.py new file mode 100644 index 0000000000..03bfbfe156 --- /dev/null +++ b/tests/test_nrrd_reader.py @@ -0,0 +1,126 @@ +# 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. + +import os +import tempfile +import unittest +from unittest.case import skipUnless + +import numpy as np +from parameterized import parameterized + +from monai.data import NrrdReader +from monai.utils.module import optional_import + +nrrd, has_nrrd = optional_import("nrrd", allow_namespace_pkg=True) + +TEST_CASE_1 = [(4, 4), "test_image.nrrd", (4, 4), np.uint8] +TEST_CASE_2 = [(4, 4, 4), "test_image.nrrd", (4, 4, 4), np.uint16] +TEST_CASE_3 = [(4, 4, 4, 4), "test_image.nrrd", (4, 4, 4, 4), np.uint32] +TEST_CASE_4 = [(1, 2, 3, 4, 5), "test_image.nrrd", (1, 2, 3, 4, 5), np.uint64] +TEST_CASE_5 = [(6, 5, 4, 3, 2, 1), "test_image.nrrd", (6, 5, 4, 3, 2, 1), np.float32] +TEST_CASE_6 = [(4,), "test_image.nrrd", (4,), np.float64] +TEST_CASE_7 = [(4, 4), ["test_image.nrrd", "test_image2.nrrd", "test_image3.nrrd"], (4, 4), np.float32] +TEST_CASE_8 = [ + (3, 4, 4, 1), + "test_image.nrrd", + (3, 4, 4, 1), + np.float32, + { + "dimension": 4, + "space": "left-posterior-superior", + "sizes": [3, 4, 4, 1], + "space directions": [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + "space origin": [0.0, 0.0, 0.0], + }, +] + + +@skipUnless(has_nrrd, "nrrd required") +class TestNrrdReader(unittest.TestCase): + def test_verify_suffix(self): + reader = NrrdReader() + self.assertFalse(reader.verify_suffix("test_image.nrd")) + reader.verify_suffix("test_image.nrrd") + reader.verify_suffix("test_image.seg.nrrd") + + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) + def test_read_int(self, data_shape, filename, expected_shape, dtype): + min_val, max_val = np.iinfo(dtype).min, np.iinfo(dtype).max + test_image = np.random.randint(min_val, max_val, size=data_shape, dtype=dtype) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, filename) + nrrd.write(filename, test_image.astype(dtype)) + reader = NrrdReader() + result = reader.read(filename) + self.assertEqual(result.array.dtype, dtype) + self.assertTupleEqual(result.array.shape, expected_shape) + self.assertTupleEqual(tuple(result.header["sizes"]), expected_shape) + np.testing.assert_allclose(result.array, test_image) + + @parameterized.expand([TEST_CASE_5, TEST_CASE_6]) + def test_read_float(self, data_shape, filename, expected_shape, dtype): + test_image = np.random.rand(*data_shape).astype(dtype) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, filename) + nrrd.write(filename, test_image.astype(dtype)) + reader = NrrdReader() + result = reader.read(filename) + self.assertEqual(result.array.dtype, dtype) + self.assertTupleEqual(result.array.shape, expected_shape) + self.assertTupleEqual(tuple(result.header["sizes"]), expected_shape) + np.testing.assert_allclose(result.array, test_image) + + @parameterized.expand([TEST_CASE_7]) + def test_read_list(self, data_shape, filenames, expected_shape, dtype): + test_image = np.random.rand(*data_shape).astype(dtype) + with tempfile.TemporaryDirectory() as tempdir: + for i, filename in enumerate(filenames): + filenames[i] = os.path.join(tempdir, filename) + nrrd.write(filenames[i], test_image.astype(dtype)) + reader = NrrdReader() + results = reader.read(filenames) + for result in results: + self.assertTupleEqual(result.array.shape, expected_shape) + self.assertTupleEqual(tuple(result.header["sizes"]), expected_shape) + np.testing.assert_allclose(result.array, test_image) + + @parameterized.expand([TEST_CASE_8]) + def test_read_with_header(self, data_shape, filename, expected_shape, dtype, reference_header): + test_image = np.random.rand(*data_shape).astype(dtype) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, filename) + nrrd.write(filename, test_image.astype(dtype), header=reference_header) + reader = NrrdReader() + image_array, image_header = reader.get_data(reader.read(filename)) + self.assertIsInstance(image_array, np.ndarray) + self.assertEqual(image_array.dtype, dtype) + self.assertTupleEqual(image_array.shape, expected_shape) + np.testing.assert_allclose(image_array, test_image) + self.assertIsInstance(image_header, dict) + self.assertTupleEqual(tuple(image_header["spatial_shape"]), expected_shape) + + @parameterized.expand([TEST_CASE_8]) + def test_read_with_header_index_order_c(self, data_shape, filename, expected_shape, dtype, reference_header): + test_image = np.random.rand(*data_shape).astype(dtype) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, filename) + nrrd.write(filename, test_image.astype(dtype), header=reference_header) + reader = NrrdReader(index_order="C") + image_array, image_header = reader.get_data(reader.read(filename)) + self.assertIsInstance(image_array, np.ndarray) + self.assertEqual(image_array.dtype, dtype) + self.assertTupleEqual(image_array.shape, expected_shape[::-1]) + self.assertTupleEqual(image_array.shape, tuple(image_header["spatial_shape"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_nuclick_transforms.py b/tests/test_nuclick_transforms.py new file mode 100644 index 0000000000..58876c3c36 --- /dev/null +++ b/tests/test_nuclick_transforms.py @@ -0,0 +1,218 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.apps.nuclick.transforms import ( + AddClickSignalsd, + AddPointGuidanceSignald, + ExtractPatchd, + FilterImaged, + FlattenLabeld, + PostFilterLabeld, + SplitLabeld, +) + +# Data Definitions +RGB_IMAGE_1 = np.array( + [[[0, 0, 0], [0, 1, 0], [0, 0, 1]], [[2, 0, 2], [0, 1, 0], [1, 0, 1]], [[3, 0, 2], [0, 1, 0], [1, 3, 1]]] +) + +LABEL_1 = np.array( + [ + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 1, 1, 1], + [1, 1, 1, 0, 1, 1, 1], + [1, 1, 1, 0, 1, 1, 1], + ], + dtype=np.uint8, +) + +LABEL_1_1 = np.array( + [ + [1, 1, 1, 0, 0, 1, 1], + [1, 1, 1, 0, 0, 1, 1], + [1, 1, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 2, 2, 2], + [1, 1, 1, 0, 2, 2, 2], + [1, 1, 1, 0, 2, 2, 2], + ], + dtype=np.uint8, +) + +LABEL_2 = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], dtype=np.uint8) + +LABEL_3 = np.array([[[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]]], dtype=np.uint8) + +LABEL_4 = np.array([[[4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4]]], dtype=np.uint8) + +IL_IMAGE_1 = np.array( + [ + [[0, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]], + [[0, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]], + [[0, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]], + ] +) + +IL_FG_IMAGE_1 = np.array([[0, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]]) + +IL_LABEL_1 = np.array( + [[[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]], dtype=np.uint8 +) + +IL_OTHERS_1 = np.array( + [[[1, 1, 1, 1, 1], [2, 0, 0, 0, 2], [3, 0, 0, 0, 3], [4, 0, 0, 0, 4], [5, 5, 5, 5, 5]]], dtype=np.uint8 +) + +IL_IMAGE_2 = np.array( + [[[0, 0, 0], [0, 1, 0], [0, 0, 1]], [[0, 0, 0], [0, 1, 0], [0, 0, 1]], [[0, 0, 0], [0, 1, 0], [0, 0, 1]]] +) + +IL_LABEL_2 = np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]]], dtype=np.uint8) + +PRED_1 = np.array( + [[[1, 1, 1, 1, 1], [2, 0, 0, 0, 2], [3, 0, 0, 0, 3], [4, 0, 0, 0, 4], [5, 5, 5, 5, 5]]], dtype=np.float32 +) + +NUC_POINTS_1 = np.array( + [ + [[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]]], + [[[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]], + ], + dtype=np.float32, +) +BB_1 = np.array([[1, 1, 3, 3], [0, 0, 2, 2]], dtype=np.uint8) + +DATA_FILTER_1 = {"image": RGB_IMAGE_1} + +DATA_FLATTEN_1 = {"label": LABEL_1} +DATA_FLATTEN_2 = {"label": LABEL_2} + +DATA_EXTRACT_1 = {"image": IL_IMAGE_1, "label": IL_LABEL_1, "centroid": (2, 2)} +DATA_EXTRACT_2 = {"image": IL_IMAGE_2, "label": IL_LABEL_2, "centroid": (1, 1)} + +DATA_SPLIT_1 = {"label": LABEL_3, "mask_value": 1} +DATA_SPLIT_2 = {"label": LABEL_4, "mask_value": 4} + +DATA_GUIDANCE_1 = {"image": IL_IMAGE_1, "label": IL_LABEL_1, "others": IL_OTHERS_1, "centroid": (2, 2)} + +DATA_CLICK_1 = {"image": IL_IMAGE_1, "foreground": [[2, 2], [1, 1]]} + +DATA_LABEL_FILTER_1 = { + "pred": PRED_1, + "nuc_points": NUC_POINTS_1, + "bounding_boxes": BB_1, + "img_height": 6, + "img_width": 6, +} + +# Result Definitions +EXTRACT_RESULT_TC1 = np.array([[[0, 0, 0], [0, 0, 0], [0, 0, 1]]], dtype=np.uint8) +EXTRACT_RESULT_TC2 = np.array([[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]], dtype=np.uint8) + +SPLIT_RESULT_TC1 = np.array([[[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]], dtype=np.uint8) +SPLIT_RESULT_TC2 = np.array([[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]], dtype=np.uint8) + +# Test Case Definitions +FILTER_IMAGE_TEST_CASE_1 = [{"keys": "image", "min_size": 1}, DATA_FILTER_1, [3, 3, 3]] + +FLATTEN_LABEL_TEST_CASE_1 = [{"keys": "label"}, DATA_FLATTEN_1, [0, 1, 2, 3]] +FLATTEN_LABEL_TEST_CASE_2 = [{"keys": "label"}, DATA_FLATTEN_2, [0]] +FLATTEN_LABEL_TEST_CASE_3 = [{"keys": "label"}, {"label": LABEL_1_1}, [0, 1, 2, 3, 4]] + +EXTRACT_TEST_CASE_1 = [{"keys": ["image", "label"], "patch_size": 3}, DATA_EXTRACT_1, [1, 3, 3]] +EXTRACT_TEST_CASE_2 = [{"keys": ["image", "label"], "patch_size": 5}, DATA_EXTRACT_1, [1, 5, 5]] +EXTRACT_TEST_CASE_3 = [{"keys": ["image", "label"], "patch_size": 1}, DATA_EXTRACT_2, [1, 1, 1]] + +EXTRACT_RESULT_TEST_CASE_1 = [{"keys": ["image", "label"], "patch_size": 3}, DATA_EXTRACT_1, EXTRACT_RESULT_TC1] +EXTRACT_RESULT_TEST_CASE_2 = [{"keys": ["image", "label"], "patch_size": 4}, DATA_EXTRACT_2, EXTRACT_RESULT_TC2] + +EXTRACT_KW_TEST_CASE_1 = [ + {"keys": ["image", "label"], "patch_size": 3, "mode": "constant"}, + DATA_EXTRACT_1, + EXTRACT_RESULT_TC1, +] + +SPLIT_TEST_CASE_1 = [{"keys": ["label"], "mask_value": "mask_value", "min_area": 1}, DATA_SPLIT_1, SPLIT_RESULT_TC1] +SPLIT_TEST_CASE_2 = [{"keys": ["label"], "mask_value": "mask_value", "min_area": 3}, DATA_SPLIT_2, SPLIT_RESULT_TC2] + +GUIDANCE_TEST_CASE_1 = [{"image": "image", "label": "label", "others": "others"}, DATA_GUIDANCE_1, [5, 5, 5]] + +CLICK_TEST_CASE_1 = [{"image": "image", "foreground": "foreground", "bb_size": 4}, DATA_CLICK_1, [2, 5, 4, 4]] + +LABEL_FILTER_TEST_CASE_1 = [{"keys": ["pred"]}, DATA_LABEL_FILTER_1, [6, 6]] + +# Test Case Classes + + +class TestFilterImaged(unittest.TestCase): + @parameterized.expand([FILTER_IMAGE_TEST_CASE_1]) + def test_correct_shape(self, arguments, input_data, expected_shape): + result = FilterImaged(**arguments)(input_data) + np.testing.assert_equal(result["image"].shape, expected_shape) + + +class TestFlattenLabeld(unittest.TestCase): + @parameterized.expand([FLATTEN_LABEL_TEST_CASE_1, FLATTEN_LABEL_TEST_CASE_2, FLATTEN_LABEL_TEST_CASE_3]) + def test_correct_num_labels(self, arguments, input_data, expected_result): + result = FlattenLabeld(**arguments)(input_data) + np.testing.assert_equal(np.unique(result["label"]), expected_result) + + +class TestExtractPatchd(unittest.TestCase): + @parameterized.expand([EXTRACT_TEST_CASE_1, EXTRACT_TEST_CASE_2, EXTRACT_TEST_CASE_3]) + def test_correct_patch_size(self, arguments, input_data, expected_shape): + result = ExtractPatchd(**arguments)(input_data) + np.testing.assert_equal(result["label"].shape, expected_shape) + + @parameterized.expand([EXTRACT_RESULT_TEST_CASE_1, EXTRACT_RESULT_TEST_CASE_2, EXTRACT_KW_TEST_CASE_1]) + def test_correct_results(self, arguments, input_data, expected_result): + result = ExtractPatchd(**arguments)(input_data) + np.testing.assert_equal(result["label"], expected_result) + + +class TestSplitLabelsd(unittest.TestCase): + @parameterized.expand([SPLIT_TEST_CASE_1, SPLIT_TEST_CASE_2]) + def test_correct_results(self, arguments, input_data, expected_result): + result = SplitLabeld(**arguments)(input_data) + np.testing.assert_equal(result["label"], expected_result) + + +class TestGuidanceSignal(unittest.TestCase): + @parameterized.expand([GUIDANCE_TEST_CASE_1]) + def test_correct_shape(self, arguments, input_data, expected_shape): + result = AddPointGuidanceSignald(**arguments)(input_data) + np.testing.assert_equal(result["image"].shape, expected_shape) + + +class TestClickSignal(unittest.TestCase): + @parameterized.expand([CLICK_TEST_CASE_1]) + def test_correct_shape(self, arguments, input_data, expected_shape): + result = AddClickSignalsd(**arguments)(input_data) + np.testing.assert_equal(result["image"].shape, expected_shape) + + +class TestPostFilterLabel(unittest.TestCase): + @parameterized.expand([LABEL_FILTER_TEST_CASE_1]) + def test_correct_shape(self, arguments, input_data, expected_shape): + result = PostFilterLabeld(**arguments)(input_data) + np.testing.assert_equal(result["pred"].shape, expected_shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_numpy_reader.py b/tests/test_numpy_reader.py index c2f3679e33..bb7686f67d 100644 --- a/tests/test_numpy_reader.py +++ b/tests/test_numpy_reader.py @@ -19,7 +19,6 @@ from monai.data import DataLoader, Dataset, NumpyReader from monai.transforms import LoadImaged -from monai.utils.enums import PostFix class TestNumpyReader(unittest.TestCase): @@ -110,8 +109,6 @@ def test_dataloader(self): num_workers=num_workers, ) for d in loader: - for s in d[PostFix.meta("image")]["spatial_shape"]: - torch.testing.assert_allclose(s, torch.as_tensor([3, 4, 5])) for c in d["image"]: torch.testing.assert_allclose(c, test_data) diff --git a/tests/test_nvtx_decorator.py b/tests/test_nvtx_decorator.py index e81c72efcf..9932b678c9 100644 --- a/tests/test_nvtx_decorator.py +++ b/tests/test_nvtx_decorator.py @@ -19,16 +19,18 @@ Compose, CuCIM, Flip, - FlipD, + Flipd, + OneOf, RandAdjustContrast, RandCuCIM, RandFlip, Randomizable, Rotate90, ToCupy, + ToNumpy, TorchVision, ToTensor, - ToTensorD, + ToTensord, ) from monai.utils import Range, optional_import from tests.utils import HAS_CUPY @@ -49,8 +51,32 @@ TEST_CASE_WRAPPER = [np.random.randn(3, 10, 10)] +TEST_CASE_RECURSIVE_0 = [ + torch.randn(3, 3), + Compose([ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()]), +] +TEST_CASE_RECURSIVE_1 = [ + torch.randn(3, 3), + Compose([ToNumpy(), Flip(), Compose([RandAdjustContrast(prob=0.0), RandFlip(prob=1.0)]), ToTensor()]), +] +TEST_CASE_RECURSIVE_2 = [ + torch.randn(3, 3), + Compose( + [ + ToNumpy(), + Flip(), + OneOf([RandAdjustContrast(prob=0.0), RandFlip(prob=1.0)], weights=[0, 1], log_stats=True), + ToTensor(), + ] + ), +] +TEST_CASE_RECURSIVE_LIST = [ + torch.randn(3, 3), + [ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()], +] + -@unittest.skipUnless(has_nvtx, "CUDA is required for NVTX Range!") +@unittest.skipUnless(has_nvtx, "Required torch._C._nvtx for NVTX Range!") class TestNVTXRangeDecorator(unittest.TestCase): @parameterized.expand([TEST_CASE_ARRAY_0, TEST_CASE_ARRAY_1]) def test_tranform_array(self, input): @@ -79,7 +105,7 @@ def test_tranform_array(self, input): @parameterized.expand([TEST_CASE_DICT_0, TEST_CASE_DICT_1]) def test_tranform_dict(self, input): - transforms = Compose([Range("random flip dict")(FlipD(keys="image")), Range()(ToTensorD("image"))]) + transforms = Compose([Range("random flip dict")(Flipd(keys="image")), Range()(ToTensord("image"))]) # Apply transforms output = transforms(input)["image"] @@ -127,6 +153,35 @@ def test_wrapper_tranforms(self, input): # Check the outputs np.testing.assert_equal(output.get(), output_r.get()) + @parameterized.expand([TEST_CASE_RECURSIVE_0, TEST_CASE_RECURSIVE_1, TEST_CASE_RECURSIVE_2]) + def test_recursive_tranforms(self, input, transforms): + transforms_range = Range(name="Recursive Compose", recursive=True)(transforms) + + # Apply transforms + output = transforms(input) + + # Apply transforms with Range + output_r = transforms_range(input) + + # Check the outputs + self.assertEqual(transforms.map_items, transforms_range.map_items) + self.assertEqual(transforms.unpack_items, transforms_range.unpack_items) + self.assertEqual(transforms.log_stats, transforms_range.log_stats) + np.testing.assert_equal(output.numpy(), output_r.numpy()) + + @parameterized.expand([TEST_CASE_RECURSIVE_LIST]) + def test_recursive_list_tranforms(self, input, transform_list): + transforms_list_range = Range(recursive=True)(transform_list) + + # Apply transforms + output = Compose(transform_list)(input) + + # Apply transforms with Range + output_r = Compose(transforms_list_range)(input) + + # Check the outputs + np.testing.assert_equal(output.numpy(), output_r.numpy()) + @parameterized.expand([TEST_CASE_ARRAY_1]) def test_tranform_randomized(self, input): # Compose deterministic and randomized transforms diff --git a/tests/test_optim_novograd.py b/tests/test_optim_novograd.py index 0cf4c35cb6..c1e63182e6 100644 --- a/tests/test_optim_novograd.py +++ b/tests/test_optim_novograd.py @@ -32,9 +32,10 @@ def build_test_cases(data): {"params": [bias], "lr": 1e-2, "amsgrad": True, "grad_averaging": True, "weight_decay": 0.1}, ] - test_cases = [] - test_cases.append([test_case_same_param, default_params, weight, bias, input]) - test_cases.append([test_case_diff_param, default_params, weight, bias, input]) + test_cases = [ + [test_case_same_param, default_params, weight, bias, input], + [test_case_diff_param, default_params, weight, bias, input], + ] return test_cases diff --git a/tests/test_orientation.py b/tests/test_orientation.py index 2b749dabad..3026305d6a 100644 --- a/tests/test_orientation.py +++ b/tests/test_orientation.py @@ -13,180 +13,227 @@ import nibabel as nib import numpy as np +import torch from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.transforms import Orientation, create_rotate, create_translate -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_DEVICES, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for device in TEST_DEVICES: TESTS.append( [ - p, {"axcodes": "RAS"}, - np.arange(12).reshape((2, 1, 2, 3)), - {"affine": np.eye(4)}, - np.arange(12).reshape((2, 1, 2, 3)), + torch.arange(12).reshape((2, 1, 2, 3)), + torch.eye(4), + torch.arange(12).reshape((2, 1, 2, 3)), "RAS", + *device, ] ) TESTS.append( [ - p, {"axcodes": "ALS"}, - np.arange(12).reshape((2, 1, 2, 3)), - {"affine": np.diag([-1, -1, 1, 1])}, - np.array([[[[3, 4, 5]], [[0, 1, 2]]], [[[9, 10, 11]], [[6, 7, 8]]]]), + torch.arange(12).reshape((2, 1, 2, 3)), + torch.as_tensor(np.diag([-1, -1, 1, 1])), + torch.tensor([[[[3, 4, 5]], [[0, 1, 2]]], [[[9, 10, 11]], [[6, 7, 8]]]]), "ALS", + *device, ] ) TESTS.append( [ - p, {"axcodes": "RAS"}, - np.arange(12).reshape((2, 1, 2, 3)), - {"affine": np.diag([-1, -1, 1, 1])}, - np.array([[[[3, 4, 5], [0, 1, 2]]], [[[9, 10, 11], [6, 7, 8]]]]), + torch.arange(12).reshape((2, 1, 2, 3)), + torch.as_tensor(np.diag([-1, -1, 1, 1])), + torch.tensor([[[[3, 4, 5], [0, 1, 2]]], [[[9, 10, 11], [6, 7, 8]]]]), "RAS", + *device, ] ) TESTS.append( [ - p, {"axcodes": "AL"}, - np.arange(6).reshape((2, 1, 3)), - {"affine": np.eye(3)}, - np.array([[[0], [1], [2]], [[3], [4], [5]]]), + torch.arange(6).reshape((2, 1, 3)), + torch.eye(3), + torch.tensor([[[0], [1], [2]], [[3], [4], [5]]]), "AL", + *device, ] ) TESTS.append( [ - p, {"axcodes": "L"}, - np.arange(6).reshape((2, 3)), - {"affine": np.eye(2)}, - np.array([[2, 1, 0], [5, 4, 3]]), + torch.arange(6).reshape((2, 3)), + torch.eye(2), + torch.tensor([[2, 1, 0], [5, 4, 3]]), "L", + *device, ] ) TESTS.append( [ - p, {"axcodes": "L"}, - np.arange(6).reshape((2, 3)), - {"affine": np.eye(2)}, - np.array([[2, 1, 0], [5, 4, 3]]), + torch.arange(6).reshape((2, 3)), + torch.eye(2), + torch.tensor([[2, 1, 0], [5, 4, 3]]), "L", + *device, ] ) TESTS.append( [ - p, {"axcodes": "L"}, - np.arange(6).reshape((2, 3)), - {"affine": np.diag([-1, 1])}, - np.arange(6).reshape((2, 3)), + torch.arange(6).reshape((2, 3)), + torch.as_tensor(np.diag([-1, 1])), + torch.arange(6).reshape((2, 3)), "L", + *device, ] ) TESTS.append( [ - p, {"axcodes": "LPS"}, - np.arange(12).reshape((2, 1, 2, 3)), - { - "affine": create_translate(3, (10, 20, 30)) + torch.arange(12).reshape((2, 1, 2, 3)), + torch.as_tensor( + create_translate(3, (10, 20, 30)) @ create_rotate(3, (np.pi / 2, np.pi / 2, np.pi / 4)) @ np.diag([-1, 1, 1, 1]) - }, - np.array([[[[2, 5]], [[1, 4]], [[0, 3]]], [[[8, 11]], [[7, 10]], [[6, 9]]]]), + ), + torch.tensor([[[[2, 5]], [[1, 4]], [[0, 3]]], [[[8, 11]], [[7, 10]], [[6, 9]]]]), "LPS", + *device, ] ) TESTS.append( [ - p, {"as_closest_canonical": True}, - np.arange(12).reshape((2, 1, 2, 3)), - { - "affine": create_translate(3, (10, 20, 30)) + torch.arange(12).reshape((2, 1, 2, 3)), + torch.as_tensor( + create_translate(3, (10, 20, 30)) @ create_rotate(3, (np.pi / 2, np.pi / 2, np.pi / 4)) @ np.diag([-1, 1, 1, 1]) - }, - np.array([[[[0, 3]], [[1, 4]], [[2, 5]]], [[[6, 9]], [[7, 10]], [[8, 11]]]]), + ), + torch.tensor([[[[0, 3]], [[1, 4]], [[2, 5]]], [[[6, 9]], [[7, 10]], [[8, 11]]]]), "RAS", + *device, ] ) TESTS.append( [ - p, {"as_closest_canonical": True}, - np.arange(6).reshape((1, 2, 3)), - {"affine": create_translate(2, (10, 20)) @ create_rotate(2, (np.pi / 3)) @ np.diag([-1, -0.2, 1])}, - np.array([[[3, 0], [4, 1], [5, 2]]]), + torch.arange(6).reshape((1, 2, 3)), + torch.as_tensor(create_translate(2, (10, 20)) @ create_rotate(2, (np.pi / 3)) @ np.diag([-1, -0.2, 1])), + torch.tensor([[[3, 0], [4, 1], [5, 2]]]), "RA", + *device, ] ) TESTS.append( [ - p, {"axcodes": "LP"}, - np.arange(6).reshape((1, 2, 3)), - {"affine": create_translate(2, (10, 20)) @ create_rotate(2, (np.pi / 3)) @ np.diag([-1, -0.2, 1])}, - np.array([[[2, 5], [1, 4], [0, 3]]]), + torch.arange(6).reshape((1, 2, 3)), + torch.as_tensor(create_translate(2, (10, 20)) @ create_rotate(2, (np.pi / 3)) @ np.diag([-1, -0.2, 1])), + torch.tensor([[[2, 5], [1, 4], [0, 3]]]), "LP", + *device, ] ) TESTS.append( [ - p, {"axcodes": "LPID", "labels": tuple(zip("LPIC", "RASD"))}, - np.zeros((1, 2, 3, 4, 5)), - {"affine": np.diag([-1, -0.2, -1, 1, 1])}, - np.zeros((1, 2, 3, 4, 5)), + torch.zeros((1, 2, 3, 4, 5)), + torch.as_tensor(np.diag([-1, -0.2, -1, 1, 1])), + torch.zeros((1, 2, 3, 4, 5)), "LPID", + *device, ] ) TESTS.append( [ - p, {"as_closest_canonical": True, "labels": tuple(zip("LPIC", "RASD"))}, - np.zeros((1, 2, 3, 4, 5)), - {"affine": np.diag([-1, -0.2, -1, 1, 1])}, - np.zeros((1, 2, 3, 4, 5)), + torch.zeros((1, 2, 3, 4, 5)), + torch.as_tensor(np.diag([-1, -0.2, -1, 1, 1])), + torch.zeros((1, 2, 3, 4, 5)), "RASD", + *device, ] ) +TESTS_TORCH = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TESTS_TORCH.append([{"axcodes": "LPS"}, torch.zeros((1, 3, 4, 5)), track_meta, *device]) + ILL_CASES = [ - # no axcodes or as_cloest_canonical - [{}, np.arange(6).reshape((2, 3)), "L"], # too short axcodes - [{"axcodes": "RA"}, np.arange(12).reshape((2, 1, 2, 3)), {"affine": np.eye(4)}], + [{"axcodes": "RA"}, torch.arange(12).reshape((2, 1, 2, 3)), torch.eye(4)] ] class TestOrientationCase(unittest.TestCase): @parameterized.expand(TESTS) - def test_ornt(self, in_type, init_param, img, data_param, expected_data, expected_code): - img = in_type(img) + def test_ornt_meta( + self, + init_param, + img: torch.Tensor, + affine: torch.Tensor, + expected_data: torch.Tensor, + expected_code: str, + device, + ): + img = MetaTensor(img, affine=affine).to(device) ornt = Orientation(**init_param) - res = ornt(img, **data_param) - if not isinstance(res, tuple): - assert_allclose(res, in_type(expected_data)) - return - assert_allclose(res[0], in_type(expected_data)) - original_affine = data_param["affine"] - np.testing.assert_allclose(original_affine, res[1]) - new_code = nib.orientations.aff2axcodes(res[2], labels=ornt.labels) + res: MetaTensor = ornt(img) + assert_allclose(res, expected_data.to(device)) + new_code = nib.orientations.aff2axcodes(res.affine.cpu(), labels=ornt.labels) self.assertEqual("".join(new_code), expected_code) + @parameterized.expand(TESTS_TORCH) + def test_ornt_torch(self, init_param, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + ornt = Orientation(**init_param) + + img = img.to(device) + expected_data = img.clone() + expected_code = ornt.axcodes + + res = ornt(img) + assert_allclose(res, expected_data) + if track_meta: + self.assertIsInstance(res, MetaTensor) + new_code = nib.orientations.aff2axcodes(res.affine.cpu(), labels=ornt.labels) + self.assertEqual("".join(new_code), expected_code) + else: + self.assertIsInstance(res, torch.Tensor) + self.assertNotIsInstance(res, MetaTensor) + @parameterized.expand(ILL_CASES) - def test_bad_params(self, init_param, img, data_param): + def test_bad_params(self, init_param, img: torch.Tensor, affine: torch.Tensor): + img = MetaTensor(img, affine=affine) with self.assertRaises(ValueError): - Orientation(**init_param)(img, **data_param) + Orientation(**init_param)(img) + + @parameterized.expand(TEST_DEVICES) + def test_inverse(self, device): + img_t = torch.rand((1, 10, 9, 8), dtype=torch.float32, device=device) + affine = torch.tensor( + [[0, 0, -1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=torch.float32, device=device + ) + meta = {"fname": "somewhere"} + img = MetaTensor(img_t, affine=affine, meta=meta) + tr = Orientation("LPS") + # check that image and affine have changed + img = tr(img) + self.assertNotEqual(img.shape, img_t.shape) + self.assertGreater((affine - img.affine).max(), 0.5) + # check that with inverse, image affine are back to how they were + img = tr.inverse(img) + self.assertEqual(img.shape, img_t.shape) + self.assertLess((affine - img.affine).max(), 1e-2) if __name__ == "__main__": diff --git a/tests/test_orientationd.py b/tests/test_orientationd.py index a4b953c8b5..1b4660a60a 100644 --- a/tests/test_orientationd.py +++ b/tests/test_orientationd.py @@ -10,94 +10,95 @@ # limitations under the License. import unittest +from typing import Optional import nibabel as nib import numpy as np +import torch +from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.transforms import Orientationd -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_DEVICES +TESTS = [] +for device in TEST_DEVICES: + TESTS.append( + [{"keys": "seg", "axcodes": "RAS"}, torch.ones((2, 1, 2, 3)), torch.eye(4), (2, 1, 2, 3), "RAS", *device] + ) + # 3d + TESTS.append( + [ + {"keys": ["img", "seg"], "axcodes": "PLI"}, + torch.ones((2, 1, 2, 3)), + torch.eye(4), + (2, 2, 1, 3), + "PLI", + *device, + ] + ) + # 2d + TESTS.append( + [{"keys": ["img", "seg"], "axcodes": "PLI"}, torch.ones((2, 1, 3)), torch.eye(4), (2, 3, 1), "PLS", *device] + ) + # 1d + TESTS.append([{"keys": ["img", "seg"], "axcodes": "L"}, torch.ones((2, 3)), torch.eye(4), (2, 3), "LAS", *device]) + # canonical + TESTS.append( + [ + {"keys": ["img", "seg"], "as_closest_canonical": True}, + torch.ones((2, 1, 2, 3)), + torch.eye(4), + (2, 1, 2, 3), + "RAS", + *device, + ] + ) -class TestOrientationdCase(unittest.TestCase): - def test_orntd(self): - data = {"seg": np.ones((2, 1, 2, 3)), PostFix.meta("seg"): {"affine": np.eye(4)}} - ornt = Orientationd(keys="seg", axcodes="RAS") - res = ornt(data) - np.testing.assert_allclose(res["seg"].shape, (2, 1, 2, 3)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("R", "A", "S")) +TESTS_TORCH = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TESTS_TORCH.append([{"keys": "seg", "axcodes": "RAS"}, torch.ones(2, 1, 2, 3), track_meta, *device]) - def test_orntd_3d(self): - for p in TEST_NDARRAYS: - data = { - "seg": p(np.ones((2, 1, 2, 3))), - "img": p(np.ones((2, 1, 2, 3))), - PostFix.meta("seg"): {"affine": np.eye(4)}, - PostFix.meta("img"): {"affine": np.eye(4)}, - } - ornt = Orientationd(keys=("img", "seg"), axcodes="PLI") - res = ornt(data) - np.testing.assert_allclose(res["img"].shape, (2, 2, 1, 3)) - np.testing.assert_allclose(res["seg"].shape, (2, 2, 1, 3)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("P", "L", "I")) - code = nib.aff2axcodes(res[PostFix.meta("img")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("P", "L", "I")) - - def test_orntd_2d(self): - data = { - "seg": np.ones((2, 1, 3)), - "img": np.ones((2, 1, 3)), - PostFix.meta("seg"): {"affine": np.eye(4)}, - PostFix.meta("img"): {"affine": np.eye(4)}, - } - ornt = Orientationd(keys=("img", "seg"), axcodes="PLI") - res = ornt(data) - np.testing.assert_allclose(res["img"].shape, (2, 3, 1)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("P", "L", "S")) - code = nib.aff2axcodes(res[PostFix.meta("img")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("P", "L", "S")) - def test_orntd_1d(self): - data = { - "seg": np.ones((2, 3)), - "img": np.ones((2, 3)), - PostFix.meta("seg"): {"affine": np.eye(4)}, - PostFix.meta("img"): {"affine": np.eye(4)}, - } - ornt = Orientationd(keys=("img", "seg"), axcodes="L") - res = ornt(data) - np.testing.assert_allclose(res["img"].shape, (2, 3)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("L", "A", "S")) - code = nib.aff2axcodes(res[PostFix.meta("img")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("L", "A", "S")) - - def test_orntd_canonical(self): - data = { - "seg": np.ones((2, 1, 2, 3)), - "img": np.ones((2, 1, 2, 3)), - PostFix.meta("seg"): {"affine": np.eye(4)}, - PostFix.meta("img"): {"affine": np.eye(4)}, - } - ornt = Orientationd(keys=("img", "seg"), as_closest_canonical=True) +class TestOrientationdCase(unittest.TestCase): + @parameterized.expand(TESTS) + def test_orntd( + self, init_param, img: torch.Tensor, affine: Optional[torch.Tensor], expected_shape, expected_code, device + ): + ornt = Orientationd(**init_param) + if affine is not None: + img = MetaTensor(img, affine=affine) + img = img.to(device) + data = {k: img.clone() for k in ornt.keys} res = ornt(data) - np.testing.assert_allclose(res["img"].shape, (2, 1, 2, 3)) - np.testing.assert_allclose(res["seg"].shape, (2, 1, 2, 3)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("R", "A", "S")) - code = nib.aff2axcodes(res[PostFix.meta("img")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("R", "A", "S")) + for k in ornt.keys: + _im = res[k] + self.assertIsInstance(_im, MetaTensor) + np.testing.assert_allclose(_im.shape, expected_shape) + code = nib.aff2axcodes(_im.affine.cpu(), ornt.ornt_transform.labels) + self.assertEqual("".join(code), expected_code) - def test_orntd_no_metadata(self): - data = {"seg": np.ones((2, 1, 2, 3))} - ornt = Orientationd(keys="seg", axcodes="RAS") + @parameterized.expand(TESTS_TORCH) + def test_orntd_torch(self, init_param, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + ornt = Orientationd(**init_param) + img = img.to(device) + expected_shape = img.shape + expected_code = ornt.ornt_transform.axcodes + data = {k: img.clone() for k in ornt.keys} res = ornt(data) - np.testing.assert_allclose(res["seg"].shape, (2, 1, 2, 3)) - code = nib.aff2axcodes(res[PostFix.meta("seg")]["affine"], ornt.ornt_transform.labels) - self.assertEqual(code, ("R", "A", "S")) + for k in ornt.keys: + _im = res[k] + np.testing.assert_allclose(_im.shape, expected_shape) + if track_meta: + self.assertIsInstance(_im, MetaTensor) + code = nib.aff2axcodes(_im.affine.cpu(), ornt.ornt_transform.labels) + self.assertEqual("".join(code), expected_code) + else: + self.assertIsInstance(_im, torch.Tensor) + self.assertNotIsInstance(_im, MetaTensor) if __name__ == "__main__": diff --git a/tests/test_pad_collation.py b/tests/test_pad_collation.py index 530e5f86a3..76a49b5b54 100644 --- a/tests/test_pad_collation.py +++ b/tests/test_pad_collation.py @@ -31,7 +31,6 @@ RandZoom, RandZoomd, ToTensor, - ToTensord, ) from monai.utils import set_determinism @@ -44,7 +43,9 @@ TESTS.append((dict, pad_collate, RandSpatialCropd("image", roi_size=[8, 7], random_size=True))) TESTS.append((dict, pad_collate, RandRotated("image", prob=1, range_x=np.pi, keep_size=False, dtype=np.float64))) TESTS.append((dict, pad_collate, RandZoomd("image", prob=1, min_zoom=1.1, max_zoom=2.0, keep_size=False))) - TESTS.append((dict, pad_collate, Compose([RandRotate90d("image", prob=1, max_k=2), ToTensord("image")]))) + TESTS.append( + (dict, pad_collate, Compose([RandRotate90d("image", prob=1, max_k=3), RandRotate90d("image", prob=1, max_k=4)])) + ) TESTS.append((list, pad_collate, RandSpatialCrop(roi_size=[8, 7], random_size=True))) TESTS.append((list, pad_collate, RandRotate(prob=1, range_x=np.pi, keep_size=False, dtype=np.float64))) @@ -97,9 +98,12 @@ def test_pad_collation(self, t_type, collate_method, transform): # check collation in forward direction for data in loader: if t_type == dict: + shapes = [] decollated_data = decollate_batch(data) for d in decollated_data: - PadListDataCollate.inverse(d) + output = PadListDataCollate.inverse(d) + shapes.append(output["image"].shape) + self.assertTrue(len(set(shapes)) > 1) # inverted shapes must be different because of random xforms if __name__ == "__main__": diff --git a/tests/test_patch_wsi_dataset_new.py b/tests/test_patch_wsi_dataset_new.py new file mode 100644 index 0000000000..fee8a03068 --- /dev/null +++ b/tests/test_patch_wsi_dataset_new.py @@ -0,0 +1,181 @@ +# 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. + +import os +import unittest +from unittest import skipUnless + +import numpy as np +from numpy.testing import assert_array_equal +from parameterized import parameterized + +from monai.data import PatchWSIDataset +from monai.data.wsi_reader import CuCIMWSIReader, OpenSlideWSIReader +from monai.utils import optional_import +from tests.utils import download_url_or_skip_test, testing_data_config + +cucim, has_cucim = optional_import("cucim") +has_cucim = has_cucim and hasattr(cucim, "CuImage") +openslide, has_osl = optional_import("openslide") +imwrite, has_tiff = optional_import("tifffile", name="imwrite") +_, has_codec = optional_import("imagecodecs") +has_tiff = has_tiff and has_codec + +FILE_KEY = "wsi_img" +FILE_URL = testing_data_config("images", FILE_KEY, "url") +base_name, extension = os.path.basename(f"{FILE_URL}"), ".tiff" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", "temp_" + base_name + extension) + +TEST_CASE_0 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "label": [1], "patch_level": 0}], "patch_size": (1, 1)}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([1])}, +] + +TEST_CASE_0_L1 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "label": [1]}], "patch_size": (1, 1), "patch_level": 1}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([1])}, +] + +TEST_CASE_0_L2 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "label": [1]}], "patch_size": (1, 1), "patch_level": 1}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([1])}, +] +TEST_CASE_1 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "patch_size": 1, "label": [1]}]}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([1])}, +] + +TEST_CASE_2 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "label": [1]}], "patch_size": 1, "patch_level": 0}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([1])}, +] + +TEST_CASE_3 = [ + {"data": [{"image": FILE_PATH, "patch_location": [0, 0], "label": [[[0, 1], [1, 0]]]}], "patch_size": 1}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([[[0, 1], [1, 0]]])}, +] + +TEST_CASE_4 = [ + { + "data": [ + {"image": FILE_PATH, "patch_location": [0, 0], "label": [[[0, 1], [1, 0]]]}, + {"image": FILE_PATH, "patch_location": [0, 0], "label": [[[1, 0], [0, 0]]]}, + ], + "patch_size": 1, + }, + [ + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([[[0, 1], [1, 0]]])}, + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([[[1, 0], [0, 0]]])}, + ], +] + +TEST_CASE_5 = [ + { + "data": [ + { + "image": FILE_PATH, + "patch_location": [0, 0], + "label": [[[0, 1], [1, 0]]], + "patch_size": 1, + "patch_level": 1, + }, + { + "image": FILE_PATH, + "patch_location": [100, 100], + "label": [[[1, 0], [0, 0]]], + "patch_size": 1, + "patch_level": 1, + }, + ] + }, + [ + {"image": np.array([[[239]], [[239]], [[239]]], dtype=np.uint8), "label": np.array([[[0, 1], [1, 0]]])}, + {"image": np.array([[[243]], [[243]], [[243]]], dtype=np.uint8), "label": np.array([[[1, 0], [0, 0]]])}, + ], +] + + +@skipUnless(has_cucim or has_osl or has_tiff, "Requires cucim, openslide, or tifffile!") +def setUpModule(): + hash_type = testing_data_config("images", FILE_KEY, "hash_type") + hash_val = testing_data_config("images", FILE_KEY, "hash_val") + download_url_or_skip_test(FILE_URL, FILE_PATH, hash_type=hash_type, hash_val=hash_val) + + +class PatchWSIDatasetTests: + class Tests(unittest.TestCase): + backend = None + + @parameterized.expand([TEST_CASE_0, TEST_CASE_0_L1, TEST_CASE_0_L2, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + def test_read_patches_str(self, input_parameters, expected): + dataset = PatchWSIDataset(reader=self.backend, **input_parameters) + sample = dataset[0] + self.assertTupleEqual(sample["label"].shape, expected["label"].shape) + self.assertTupleEqual(sample["image"].shape, expected["image"].shape) + self.assertIsNone(assert_array_equal(sample["label"], expected["label"])) + self.assertIsNone(assert_array_equal(sample["image"], expected["image"])) + + @parameterized.expand([TEST_CASE_0, TEST_CASE_0_L1, TEST_CASE_0_L2, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + def test_read_patches_class(self, input_parameters, expected): + if self.backend == "openslide": + reader = OpenSlideWSIReader + elif self.backend == "cucim": + reader = CuCIMWSIReader + else: + raise ValueError("Unsupported backend: {self.backend}") + dataset = PatchWSIDataset(reader=reader, **input_parameters) + sample = dataset[0] + self.assertTupleEqual(sample["label"].shape, expected["label"].shape) + self.assertTupleEqual(sample["image"].shape, expected["image"].shape) + self.assertIsNone(assert_array_equal(sample["label"], expected["label"])) + self.assertIsNone(assert_array_equal(sample["image"], expected["image"])) + + @parameterized.expand([TEST_CASE_0, TEST_CASE_0_L1, TEST_CASE_0_L2, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + def test_read_patches_object(self, input_parameters, expected): + if self.backend == "openslide": + reader = OpenSlideWSIReader(level=input_parameters.get("patch_level", 0)) + elif self.backend == "cucim": + reader = CuCIMWSIReader(level=input_parameters.get("patch_level", 0)) + else: + raise ValueError("Unsupported backend: {self.backend}") + dataset = PatchWSIDataset(reader=reader, **input_parameters) + sample = dataset[0] + self.assertTupleEqual(sample["label"].shape, expected["label"].shape) + self.assertTupleEqual(sample["image"].shape, expected["image"].shape) + self.assertIsNone(assert_array_equal(sample["label"], expected["label"])) + self.assertIsNone(assert_array_equal(sample["image"], expected["image"])) + + @parameterized.expand([TEST_CASE_4, TEST_CASE_5]) + def test_read_patches_str_multi(self, input_parameters, expected): + dataset = PatchWSIDataset(reader=self.backend, **input_parameters) + for i in range(len(dataset)): + self.assertTupleEqual(dataset[i]["label"].shape, expected[i]["label"].shape) + self.assertTupleEqual(dataset[i]["image"].shape, expected[i]["image"].shape) + self.assertIsNone(assert_array_equal(dataset[i]["label"], expected[i]["label"])) + self.assertIsNone(assert_array_equal(dataset[i]["image"], expected[i]["image"])) + + +@skipUnless(has_cucim, "Requires cucim") +class TestPatchWSIDatasetCuCIM(PatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "cucim" + + +@skipUnless(has_osl, "Requires openslide") +class TestPatchWSIDatasetOpenSlide(PatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "openslide" + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_patchembedding.py b/tests/test_patchembedding.py index 4af2b47ba5..6971eb0463 100644 --- a/tests/test_patchembedding.py +++ b/tests/test_patchembedding.py @@ -13,10 +13,11 @@ from unittest import skipUnless import torch +import torch.nn as nn from parameterized import parameterized from monai.networks import eval_mode -from monai.networks.blocks.patchembedding import PatchEmbeddingBlock +from monai.networks.blocks.patchembedding import PatchEmbed, PatchEmbeddingBlock from monai.utils import optional_import einops, has_einops = optional_import("einops") @@ -48,6 +49,26 @@ test_case[0]["spatial_dims"] = 2 # type: ignore TEST_CASE_PATCHEMBEDDINGBLOCK.append(test_case) +TEST_CASE_PATCHEMBED = [] +for patch_size in [2]: + for in_chans in [1, 4]: + for img_size in [96]: + for embed_dim in [6, 12]: + for norm_layer in [nn.LayerNorm]: + for nd in [2, 3]: + test_case = [ + { + "patch_size": (patch_size,) * nd, + "in_chans": in_chans, + "embed_dim": embed_dim, + "norm_layer": norm_layer, + "spatial_dims": nd, + }, + (2, in_chans, *([img_size] * nd)), + (2, embed_dim, *([img_size // patch_size] * nd)), + ] + TEST_CASE_PATCHEMBED.append(test_case) + class TestPatchEmbeddingBlock(unittest.TestCase): @parameterized.expand(TEST_CASE_PATCHEMBEDDINGBLOCK) @@ -115,5 +136,19 @@ def test_ill_arg(self): ) +class TestPatchEmbed(unittest.TestCase): + @parameterized.expand(TEST_CASE_PATCHEMBED) + @skipUnless(has_einops, "Requires einops") + def test_shape(self, input_param, input_shape, expected_shape): + net = PatchEmbed(**input_param) + with eval_mode(net): + result = net(torch.randn(input_shape)) + self.assertEqual(result.shape, expected_shape) + + def test_ill_arg(self): + with self.assertRaises(ValueError): + PatchEmbed(patch_size=(2, 2, 2), in_chans=1, embed_dim=24, norm_layer=nn.LayerNorm, spatial_dims=5) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_plot_2d_or_3d_image.py b/tests/test_plot_2d_or_3d_image.py index 2e4adb93e3..5325c6a294 100644 --- a/tests/test_plot_2d_or_3d_image.py +++ b/tests/test_plot_2d_or_3d_image.py @@ -15,12 +15,13 @@ import torch from parameterized import parameterized -from torch.utils.tensorboard import SummaryWriter from monai.utils import optional_import from monai.visualize import plot_2d_or_3d_image from tests.utils import SkipIfNoModule +SummaryWriter, has_tb = optional_import("torch.utils.tensorboard", name="SummaryWriter") + SummaryWriterX, _ = optional_import("tensorboardX", name="SummaryWriter") TEST_CASE_1 = [(1, 1, 10, 10)] @@ -34,6 +35,7 @@ TEST_CASE_5 = [(1, 3, 10, 10, 10)] +@unittest.skipUnless(has_tb, "Requires SummaryWriter installation") class TestPlot2dOr3dImage(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5]) def test_tb_image(self, shape): diff --git a/tests/test_print_transform_backends.py b/tests/test_print_transform_backends.py index 2db00fea39..e714003769 100644 --- a/tests/test_print_transform_backends.py +++ b/tests/test_print_transform_backends.py @@ -22,6 +22,5 @@ def test_get_number_of_conversions(self): if __name__ == "__main__": - # unittest.main() a = TestPrintTransformBackends() a.test_get_number_of_conversions() diff --git a/tests/test_rand_adjust_contrast.py b/tests/test_rand_adjust_contrast.py index eaeff70d51..5dc800793e 100644 --- a/tests/test_rand_adjust_contrast.py +++ b/tests/test_rand_adjust_contrast.py @@ -27,7 +27,8 @@ class TestRandAdjustContrast(NumpyImageTestCase2D): def test_correct_results(self, gamma): adjuster = RandAdjustContrast(prob=1.0, gamma=gamma) for p in TEST_NDARRAYS: - result = adjuster(p(self.imt)) + im = p(self.imt) + result = adjuster(im) epsilon = 1e-7 img_min = self.imt.min() img_range = self.imt.max() - img_min @@ -35,7 +36,7 @@ def test_correct_results(self, gamma): np.power(((self.imt - img_min) / float(img_range + epsilon)), adjuster.gamma_value) * img_range + img_min ) - assert_allclose(expected, result, rtol=1e-05, type_test=False) + assert_allclose(result, expected, rtol=1e-05, type_test=False) if __name__ == "__main__": diff --git a/tests/test_rand_adjust_contrastd.py b/tests/test_rand_adjust_contrastd.py index e5f1f6099a..b355ac3e4f 100644 --- a/tests/test_rand_adjust_contrastd.py +++ b/tests/test_rand_adjust_contrastd.py @@ -35,7 +35,7 @@ def test_correct_results(self, gamma): np.power(((self.imt - img_min) / float(img_range + epsilon)), adjuster.adjuster.gamma_value) * img_range + img_min ) - assert_allclose(expected, result["img"], rtol=1e-05, type_test=False) + assert_allclose(result["img"], expected, rtol=1e-05, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_affine.py b/tests/test_rand_affine.py index 8de408ab84..b5bc67ffb1 100644 --- a/tests/test_rand_affine.py +++ b/tests/test_rand_affine.py @@ -17,12 +17,12 @@ from monai.transforms import RandAffine from monai.utils.type_conversion import convert_data_type -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, is_tf32_env _rtol = 1e-3 if is_tf32_env() else 1e-4 TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [dict(device=device), {"img": p(torch.arange(27).reshape((3, 3, 3)))}, p(np.arange(27).reshape((3, 3, 3)))] @@ -126,10 +126,15 @@ ) TEST_CASES_SKIPPED_CONSISTENCY = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for in_dtype in (np.int32, np.float32): TEST_CASES_SKIPPED_CONSISTENCY.append((p(np.arange(9 * 10).reshape(1, 9, 10)), in_dtype)) +TEST_RANDOMIZE = [] +for cache_grid in (False, True): + for initial_randomize in (False, True): + TEST_RANDOMIZE.append((initial_randomize, cache_grid)) + class TestRandAffine(unittest.TestCase): @parameterized.expand(TESTS) @@ -139,7 +144,7 @@ def test_rand_affine(self, input_param, input_data, expected_val): result = g(**input_data) if input_param.get("cache_grid", False): self.assertTrue(g._cached_grid is not None) - assert_allclose(result, expected_val, rtol=_rtol, atol=1e-4) + assert_allclose(result, expected_val, rtol=_rtol, atol=1e-4, type_test="tensor") def test_ill_cache(self): with self.assertWarns(UserWarning): @@ -162,6 +167,31 @@ def test_skipped_transform_consistency(self, im, in_dtype): # check matching dtype self.assertEqual(out1.dtype, out2.dtype) + @parameterized.expand(TEST_RANDOMIZE) + def test_no_randomize(self, initial_randomize, cache_grid): + rand_affine = RandAffine( + prob=1, + rotate_range=(np.pi / 6, 0, 0), + translate_range=((-2, 2), (-2, 2), (-2, 2)), + scale_range=((-0.1, 0.1), (-0.1, 0.1), (-0.1, 0.1)), + spatial_size=(16, 16, 16), + cache_grid=cache_grid, + padding_mode="zeros", + ) + if initial_randomize: + rand_affine.randomize(None) + + arr = torch.randn((1, 16, 16, 16)) * 100 + + arr1 = rand_affine(arr, randomize=False) + m1 = rand_affine.rand_affine_grid.get_transformation_matrix() + + arr2 = rand_affine(arr, randomize=False) + m2 = rand_affine.rand_affine_grid.get_transformation_matrix() + + assert_allclose(m1, m2) + assert_allclose(arr1, arr2) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_rand_affine_grid.py b/tests/test_rand_affine_grid.py index 60ac40f468..6a40d39e4e 100644 --- a/tests/test_rand_affine_grid.py +++ b/tests/test_rand_affine_grid.py @@ -16,12 +16,12 @@ from parameterized import parameterized from monai.transforms import RandAffineGrid -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, is_tf32_env -_rtol = 1e-1 if is_tf32_env else 1e-4 +_rtol = 1e-1 if is_tf32_env() else 1e-4 TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append([{"device": device}, {"grid": p(torch.ones((3, 3, 3)))}, p(np.ones((3, 3, 3)))]) TESTS.append( diff --git a/tests/test_rand_affined.py b/tests/test_rand_affined.py index 882b5554e6..a33496895c 100644 --- a/tests/test_rand_affined.py +++ b/tests/test_rand_affined.py @@ -9,214 +9,243 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools import unittest import numpy as np import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandAffined from monai.utils import GridSampleMode -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import assert_allclose, is_tf32_env _rtol = 1e-3 if is_tf32_env() else 1e-4 TESTS = [] -for p in TEST_NDARRAYS: - for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: - TESTS.append( - [ - dict(device=device, spatial_size=None, keys=("img", "seg")), - {"img": p(torch.arange(27).reshape((3, 3, 3))), "seg": p(torch.arange(27).reshape((3, 3, 3)))}, - p(np.arange(27).reshape((3, 3, 3))), - ] - ) - TESTS.append( - [ - dict(device=device, spatial_size=(2, 2), keys=("img", "seg")), - {"img": p(torch.ones((3, 3, 3))), "seg": p(torch.ones((3, 3, 3)))}, - p(np.ones((3, 2, 2))), - ] - ) - TESTS.append( - [ - dict(device=device, spatial_size=(2, 2), cache_grid=True, keys=("img", "seg")), - {"img": p(torch.ones((3, 3, 3))), "seg": p(torch.ones((3, 3, 3)))}, - p(np.ones((3, 2, 2))), - ] - ) - TESTS.append( - [ - dict(device=device, spatial_size=(2, 2, 2), keys=("img", "seg")), - {"img": p(torch.ones((1, 3, 3, 3))), "seg": p(torch.ones((1, 3, 3, 3)))}, - p(torch.ones((1, 2, 2, 2))), - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - spatial_size=(2, 2, 2), - padding_mode="zeros", - device=device, - keys=("img", "seg"), - mode="bilinear", - ), - {"img": p(torch.ones((1, 3, 3, 3))), "seg": p(torch.ones((1, 3, 3, 3)))}, - p(torch.tensor([[[[0.3658, 1.0000], [1.0000, 1.0000]], [[1.0000, 1.0000], [1.0000, 0.9333]]]])), - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - scale_range=[0.1, 0.2], - spatial_size=(3, 3), - keys=("img", "seg"), - device=device, - ), - {"img": p(torch.arange(64).reshape((1, 8, 8))), "seg": p(torch.arange(64).reshape((1, 8, 8)))}, - p( + +for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: + TESTS.append( + [ + dict(device=device, spatial_size=None, keys=("img", "seg")), + { + "img": MetaTensor(torch.arange(27).reshape((3, 3, 3))), + "seg": MetaTensor(torch.arange(27).reshape((3, 3, 3))), + }, + torch.arange(27).reshape((3, 3, 3)), + ] + ) + TESTS.append( + [ + dict(device=device, spatial_size=(2, 2), keys=("img", "seg")), + {"img": MetaTensor(torch.ones((3, 3, 3))), "seg": MetaTensor(torch.ones((3, 3, 3)))}, + torch.ones((3, 2, 2)), + ] + ) + TESTS.append( + [ + dict(device=device, spatial_size=(2, 2), cache_grid=True, keys=("img", "seg")), + {"img": MetaTensor(torch.ones((3, 3, 3))), "seg": MetaTensor(torch.ones((3, 3, 3)))}, + torch.ones((3, 2, 2)), + ] + ) + TESTS.append( + [ + dict(device=device, spatial_size=(2, 2, 2), keys=("img", "seg")), + {"img": MetaTensor(torch.ones((1, 3, 3, 3))), "seg": MetaTensor(torch.ones((1, 3, 3, 3)))}, + torch.ones((1, 2, 2, 2)), + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + spatial_size=(2, 2, 2), + padding_mode="zeros", + device=device, + keys=("img", "seg"), + mode="bilinear", + ), + {"img": MetaTensor(torch.ones((1, 3, 3, 3))), "seg": MetaTensor(torch.ones((1, 3, 3, 3)))}, + torch.tensor([[[[0.3658, 1.0000], [1.0000, 1.0000]], [[1.0000, 1.0000], [1.0000, 0.9333]]]]), + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + scale_range=[0.1, 0.2], + spatial_size=(3, 3), + keys=("img", "seg"), + device=device, + ), + { + "img": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + "seg": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + }, + torch.tensor([[[18.7362, 15.5820, 12.4278], [27.3988, 24.2446, 21.0904], [36.0614, 32.9072, 29.7530]]]), + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + mode=("bilinear", "nearest"), + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + scale_range=[0.1, 0.2], + spatial_size=(3, 3), + keys=("img", "seg"), + device=device, + ), + { + "img": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + "seg": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + }, + { + "img": MetaTensor( torch.tensor( - [[[18.7362, 15.5820, 12.4278], [27.3988, 24.2446, 21.0904], [36.0614, 32.9072, 29.7530]]] - ) - ), - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - mode=("bilinear", "nearest"), - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - scale_range=[0.1, 0.2], - spatial_size=(3, 3), - keys=("img", "seg"), - device=device, - ), - {"img": p(torch.arange(64).reshape((1, 8, 8))), "seg": p(torch.arange(64).reshape((1, 8, 8)))}, - { - "img": p( - np.array( + [ [ - [ - [18.736153, 15.581954, 12.4277525], - [27.398798, 24.244598, 21.090399], - [36.061443, 32.90724, 29.753046], - ] + [18.736153, 15.581954, 12.4277525], + [27.398798, 24.244598, 21.090399], + [36.061443, 32.90724, 29.753046], ] - ) - ), - "seg": p(np.array([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), - }, - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - spatial_size=(2, 2, 2), - padding_mode="zeros", - device=device, - keys=("img", "seg"), - mode=GridSampleMode.BILINEAR, - ), - {"img": p(torch.ones((1, 3, 3, 3))), "seg": p(torch.ones((1, 3, 3, 3)))}, - p(torch.tensor([[[[0.3658, 1.0000], [1.0000, 1.0000]], [[1.0000, 1.0000], [1.0000, 0.9333]]]])), - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - mode=(GridSampleMode.BILINEAR, GridSampleMode.NEAREST), - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - scale_range=[0.1, 0.2], - spatial_size=(3, 3), - keys=("img", "seg"), - device=device, + ] + ) ), - {"img": p(torch.arange(64).reshape((1, 8, 8))), "seg": p(torch.arange(64).reshape((1, 8, 8)))}, - { - "img": p( - np.array( + "seg": MetaTensor(torch.tensor([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), + }, + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + spatial_size=(2, 2, 2), + padding_mode="zeros", + device=device, + keys=("img", "seg"), + mode=GridSampleMode.BILINEAR, + ), + {"img": MetaTensor(torch.ones((1, 3, 3, 3))), "seg": MetaTensor(torch.ones((1, 3, 3, 3)))}, + torch.tensor([[[[0.3658, 1.0000], [1.0000, 1.0000]], [[1.0000, 1.0000], [1.0000, 0.9333]]]]), + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + mode=(GridSampleMode.BILINEAR, GridSampleMode.NEAREST), + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + scale_range=[0.1, 0.2], + spatial_size=(3, 3), + keys=("img", "seg"), + device=device, + ), + { + "img": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + "seg": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + }, + { + "img": MetaTensor( + np.array( + [ [ - [ - [18.736153, 15.581954, 12.4277525], - [27.398798, 24.244598, 21.090399], - [36.061443, 32.90724, 29.753046], - ] + [18.736153, 15.581954, 12.4277525], + [27.398798, 24.244598, 21.090399], + [36.061443, 32.90724, 29.753046], ] - ) - ), - "seg": p(np.array([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), - }, - ] - ) - TESTS.append( - [ - dict( - prob=0.9, - mode=(GridSampleMode.BILINEAR, GridSampleMode.NEAREST), - rotate_range=(np.pi / 2,), - shear_range=[1, 2], - translate_range=[2, 1], - scale_range=[0.1, 0.2], - spatial_size=(3, 3), - cache_grid=True, - keys=("img", "seg"), - device=device, + ] + ) ), - {"img": p(torch.arange(64).reshape((1, 8, 8))), "seg": p(torch.arange(64).reshape((1, 8, 8)))}, - { - "img": p( - np.array( + "seg": MetaTensor(np.array([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), + }, + ] + ) + TESTS.append( + [ + dict( + prob=0.9, + mode=(GridSampleMode.BILINEAR, GridSampleMode.NEAREST), + rotate_range=(np.pi / 2,), + shear_range=[1, 2], + translate_range=[2, 1], + scale_range=[0.1, 0.2], + spatial_size=(3, 3), + cache_grid=True, + keys=("img", "seg"), + device=device, + ), + { + "img": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + "seg": MetaTensor(torch.arange(64).reshape((1, 8, 8))), + }, + { + "img": MetaTensor( + torch.tensor( + [ [ - [ - [18.736153, 15.581954, 12.4277525], - [27.398798, 24.244598, 21.090399], - [36.061443, 32.90724, 29.753046], - ] + [18.736153, 15.581954, 12.4277525], + [27.398798, 24.244598, 21.090399], + [36.061443, 32.90724, 29.753046], ] - ) - ), - "seg": p(np.array([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), - }, - ] - ) + ] + ) + ), + "seg": MetaTensor(torch.tensor([[[19.0, 20.0, 12.0], [27.0, 28.0, 20.0], [35.0, 36.0, 29.0]]])), + }, + ] + ) class TestRandAffined(unittest.TestCase): - @parameterized.expand(TESTS) - def test_rand_affined(self, input_param, input_data, expected_val): + @parameterized.expand(x + [y] for x, y in itertools.product(TESTS, (False, True))) + def test_rand_affined(self, input_param, input_data, expected_val, track_meta): + set_track_meta(track_meta) g = RandAffined(**input_param).set_random_state(123) res = g(input_data) if input_param.get("cache_grid", False): self.assertTrue(g.rand_affine._cached_grid is not None) for key in res: result = res[key] - if "_transforms" in key: - continue + if track_meta: + self.assertIsInstance(result, MetaTensor) + self.assertEqual(len(result.applied_operations), 1) expected = expected_val[key] if isinstance(expected_val, dict) else expected_val - assert_allclose(result, expected, rtol=_rtol, atol=1e-3) + assert_allclose(result, expected, rtol=_rtol, atol=1e-3, type_test=False) g.set_random_state(4) res = g(input_data) + if not track_meta: + return + # affine should be tensor because the resampler only supports pytorch backend - self.assertTrue(isinstance(res["img_transforms"][0]["extra_info"]["affine"], torch.Tensor)) + if isinstance(res["img"], MetaTensor) and "extra_info" in res["img"].applied_operations[0]: + if not res["img"].applied_operations[-1]["extra_info"]["do_resampling"]: + return + affine_img = res["img"].applied_operations[0]["extra_info"]["rand_affine_info"]["extra_info"]["affine"] + affine_seg = res["seg"].applied_operations[0]["extra_info"]["rand_affine_info"]["extra_info"]["affine"] + assert_allclose(affine_img, affine_seg, rtol=_rtol, atol=1e-3) + + res_inv = g.inverse(res) + for k, v in res_inv.items(): + self.assertIsInstance(v, MetaTensor) + self.assertEqual(len(v.applied_operations), 0) + self.assertTupleEqual(v.shape, input_data[k].shape) def test_ill_cache(self): with self.assertWarns(UserWarning): diff --git a/tests/test_rand_axis_flip.py b/tests/test_rand_axis_flip.py index b7c504557f..7458b9d6dd 100644 --- a/tests/test_rand_axis_flip.py +++ b/tests/test_rand_axis_flip.py @@ -12,18 +12,28 @@ import unittest import numpy as np +import torch +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandAxisFlip -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion class TestRandAxisFlip(NumpyImageTestCase2D): def test_correct_results(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: flip = RandAxisFlip(prob=1.0) - result = flip(p(self.imt[0])) + im = p(self.imt[0]) + result = flip(im) expected = [np.flip(channel, flip._axis) for channel in self.imt[0]] - assert_allclose(result, p(np.stack(expected))) + assert_allclose(result, p(np.stack(expected)), type_test="tensor") + test_local_inversion(flip, result, im) + + set_track_meta(False) + result = flip(im) + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_rand_axis_flipd.py b/tests/test_rand_axis_flipd.py index ff97d5dc1e..a62da88af3 100644 --- a/tests/test_rand_axis_flipd.py +++ b/tests/test_rand_axis_flipd.py @@ -12,19 +12,28 @@ import unittest import numpy as np +import torch +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandAxisFlipd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase3D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase3D, assert_allclose, test_local_inversion class TestRandAxisFlip(NumpyImageTestCase3D): def test_correct_results(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: flip = RandAxisFlipd(keys="img", prob=1.0) - result = flip({"img": p(self.imt[0])})["img"] - + im = p(self.imt[0]) + result = flip({"img": im}) + test_local_inversion(flip, result, {"img": im}, "img") expected = [np.flip(channel, flip.flipper._axis) for channel in self.imt[0]] - assert_allclose(result, p(np.stack(expected))) + assert_allclose(result["img"], p(np.stack(expected)), type_test="tensor") + + set_track_meta(False) + result = flip({"img": im})["img"] + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_rand_bias_field.py b/tests/test_rand_bias_field.py index b3aa8e9174..690c4022eb 100644 --- a/tests/test_rand_bias_field.py +++ b/tests/test_rand_bias_field.py @@ -16,6 +16,7 @@ from parameterized import parameterized from monai.transforms import RandBiasField +from tests.utils import TEST_NDARRAYS TEST_CASES_2D = [{"prob": 1.0}, (3, 32, 32)] TEST_CASES_3D = [{"prob": 1.0}, (3, 32, 32, 32)] @@ -29,10 +30,10 @@ class TestRandBiasField(unittest.TestCase): @parameterized.expand([TEST_CASES_2D, TEST_CASES_3D]) def test_output_shape(self, class_args, img_shape): - for fn in (np.random, torch): + for p in TEST_NDARRAYS: for degree in [1, 2, 3]: bias_field = RandBiasField(degree=degree, **class_args) - img = fn.rand(*img_shape) + img = p(np.random.rand(*img_shape)) output = bias_field(img) np.testing.assert_equal(output.shape, img_shape) self.assertTrue(output.dtype in (np.float32, torch.float32)) diff --git a/tests/test_rand_bias_fieldd.py b/tests/test_rand_bias_fieldd.py index da08cfe053..05a5a1b636 100644 --- a/tests/test_rand_bias_fieldd.py +++ b/tests/test_rand_bias_fieldd.py @@ -33,7 +33,6 @@ def test_output_shape(self, class_args, img_shape): img = np.random.rand(*img_shape) output = bias_field({key: img}) np.testing.assert_equal(output[key].shape, img_shape) - np.testing.assert_equal(output[key].dtype, bias_field.rand_bias_field.dtype) @parameterized.expand([TEST_CASES_2D_ZERO_RANGE]) def test_zero_range(self, class_args, img_shape): diff --git a/tests/test_rand_coarse_dropout.py b/tests/test_rand_coarse_dropout.py index a05d323277..cc05edbf02 100644 --- a/tests/test_rand_coarse_dropout.py +++ b/tests/test_rand_coarse_dropout.py @@ -17,6 +17,7 @@ from monai.transforms import RandCoarseDropout from monai.utils import fall_back_tuple +from tests.utils import TEST_NDARRAYS, assert_allclose TEST_CASE_0 = [ {"holes": 2, "spatial_size": [2, 2, 2], "fill_value": 5, "prob": 1.0}, @@ -64,9 +65,10 @@ class TestRandCoarseDropout(unittest.TestCase): [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7] ) def test_value(self, input_param, input_data): - dropout = RandCoarseDropout(**input_param) - result = dropout(input_data) - self.assertEqual(type(result), type(input_data)) + for p in TEST_NDARRAYS: + dropout = RandCoarseDropout(**input_param) + im = p(input_data) + result = dropout(im) holes = input_param.get("holes") max_holes = input_param.get("max_holes") spatial_size = fall_back_tuple(input_param.get("spatial_size"), input_data.shape[1:]) @@ -84,7 +86,7 @@ def test_value(self, input_param, input_data): if input_param.get("dropout_holes", True): fill_value = input_param.get("fill_value", None) if isinstance(fill_value, (int, float)): - np.testing.assert_allclose(data, fill_value) + assert_allclose(data, fill_value, type_test=False) elif fill_value is not None: min_value = data.min() max_value = data.max() @@ -92,7 +94,7 @@ def test_value(self, input_param, input_data): self.assertGreaterEqual(min_value, fill_value[0]) self.assertLess(max_value, fill_value[1]) else: - np.testing.assert_allclose(data, input_data[h]) + assert_allclose(data, input_data[h], type_test=False) if max_spatial_size is None: self.assertTupleEqual(data.shape[1:], tuple(spatial_size)) diff --git a/tests/test_rand_crop_by_label_classes.py b/tests/test_rand_crop_by_label_classes.py index 11d73df74e..b1165e8986 100644 --- a/tests/test_rand_crop_by_label_classes.py +++ b/tests/test_rand_crop_by_label_classes.py @@ -15,10 +15,10 @@ from parameterized import parameterized from monai.transforms import ClassesToIndices, RandCropByLabelClasses -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL TESTS_INDICES, TESTS_SHAPE = [], [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: # One-Hot label TESTS_INDICES.append( [ diff --git a/tests/test_rand_crop_by_label_classesd.py b/tests/test_rand_crop_by_label_classesd.py index 92780458e0..9a99ebab29 100644 --- a/tests/test_rand_crop_by_label_classesd.py +++ b/tests/test_rand_crop_by_label_classesd.py @@ -15,10 +15,10 @@ from parameterized import parameterized from monai.transforms import ClassesToIndicesd, RandCropByLabelClassesd -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TESTS.append( [ # One-Hot label diff --git a/tests/test_rand_crop_by_pos_neg_label.py b/tests/test_rand_crop_by_pos_neg_label.py index f8b8a77a45..f6da393ab9 100644 --- a/tests/test_rand_crop_by_pos_neg_label.py +++ b/tests/test_rand_crop_by_pos_neg_label.py @@ -16,10 +16,9 @@ from parameterized import parameterized from monai.transforms import RandCropByPosNegLabel -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL -TESTS = [] -TESTS.append( +TESTS = [ [ { "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), @@ -32,9 +31,7 @@ }, {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, (3, 2, 2, 3), - ] -) -TESTS.append( + ], [ { "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), @@ -47,9 +44,7 @@ }, {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, (3, 2, 2, 2), - ] -) -TESTS.append( + ], [ { "label": None, @@ -66,9 +61,7 @@ "image": np.random.randint(0, 2, size=[3, 3, 3, 3]), }, (3, 2, 2, 2), - ] -) -TESTS.append( + ], [ { "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), @@ -81,9 +74,7 @@ }, {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, (3, 3, 3, 2), - ] -) -TESTS.append( + ], [ { "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), @@ -96,8 +87,8 @@ }, {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, (3, 3, 3, 3), - ] -) + ], +] class TestRandCropByPosNegLabel(unittest.TestCase): @@ -112,12 +103,13 @@ def convert_data_type(im_type, d, keys=("img", "image", "label")): @parameterized.expand(TESTS) def test_type_shape(self, input_param, input_data, expected_shape): results = [] - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: input_param_mod = self.convert_data_type(p, input_param) input_data_mod = self.convert_data_type(p, input_data) cropper = RandCropByPosNegLabel(**input_param_mod) cropper.set_random_state(0) result = cropper(**input_data_mod) + self.assertListEqual(cropper.spatial_size, input_param["spatial_size"]) self.assertIsInstance(result, list) self.assertTupleEqual(result[0].shape, expected_shape) diff --git a/tests/test_rand_crop_by_pos_neg_labeld.py b/tests/test_rand_crop_by_pos_neg_labeld.py index df85b29b00..64673bf4bf 100644 --- a/tests/test_rand_crop_by_pos_neg_labeld.py +++ b/tests/test_rand_crop_by_pos_neg_labeld.py @@ -16,8 +16,7 @@ from parameterized import parameterized from monai.transforms import RandCropByPosNegLabeld -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL TESTS = [ [ @@ -35,7 +34,6 @@ "image": np.random.randint(0, 2, size=[3, 3, 3, 3]), "extra": np.random.randint(0, 2, size=[3, 3, 3, 3]), "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), - PostFix.meta("image"): {"affine": np.eye(3), "shape": "CHWD"}, }, (3, 3, 2, 2), ], @@ -54,7 +52,6 @@ "image": np.random.randint(0, 2, size=[3, 3, 3, 3]), "extra": np.random.randint(0, 2, size=[3, 3, 3, 3]), "label": np.random.randint(0, 2, size=[3, 3, 3, 3]), - PostFix.meta("label"): {"affine": np.eye(3), "shape": "CHWD"}, }, (3, 2, 2, 2), ], @@ -69,12 +66,7 @@ "image_key": None, "image_threshold": 0, }, - { - "image": np.zeros([3, 3, 3, 3]) - 1, - "extra": np.zeros([3, 3, 3, 3]), - "label": np.ones([3, 3, 3, 3]), - PostFix.meta("extra"): {"affine": np.eye(3), "shape": "CHWD"}, - }, + {"image": np.zeros([3, 3, 3, 3]) - 1, "extra": np.zeros([3, 3, 3, 3]), "label": np.ones([3, 3, 3, 3])}, (3, 2, 2, 2), ], [ @@ -89,12 +81,7 @@ "image_threshold": 0, "allow_smaller": True, }, - { - "image": np.zeros([3, 3, 3, 3]) - 1, - "extra": np.zeros([3, 3, 3, 3]), - "label": np.ones([3, 3, 3, 3]), - PostFix.meta("extra"): {"affine": np.eye(3), "shape": "CHWD"}, - }, + {"image": np.zeros([3, 3, 3, 3]) - 1, "extra": np.zeros([3, 3, 3, 3]), "label": np.ones([3, 3, 3, 3])}, (3, 3, 3, 2), ], [ @@ -109,12 +96,7 @@ "image_threshold": 0, "allow_smaller": True, }, - { - "image": np.zeros([3, 3, 3, 3]) - 1, - "extra": np.zeros([3, 3, 3, 3]), - "label": np.ones([3, 3, 3, 3]), - PostFix.meta("extra"): {"affine": np.eye(3), "shape": "CHWD"}, - }, + {"image": np.zeros([3, 3, 3, 3]) - 1, "extra": np.zeros([3, 3, 3, 3]), "label": np.ones([3, 3, 3, 3])}, (3, 3, 3, 3), ], ] @@ -131,12 +113,13 @@ def convert_data_type(im_type, d, keys=("img", "image", "label")): @parameterized.expand(TESTS) def test_type_shape(self, input_param, input_data, expected_shape): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: input_param_mod = self.convert_data_type(p, input_param) input_data_mod = self.convert_data_type(p, input_data) cropper = RandCropByPosNegLabeld(**input_param_mod) cropper.set_random_state(0) result = cropper(input_data_mod) + self.assertListEqual(cropper.cropper.spatial_size, input_param["spatial_size"]) self.assertIsInstance(result, list) @@ -145,7 +128,14 @@ def test_type_shape(self, input_param, input_data, expected_shape): for k in ("image", "extra", "label"): self.assertTupleEqual(result[0][k].shape, expected_shape) for i, item in enumerate(result): - self.assertEqual(item[PostFix.meta(k)]["patch_index"], i) + self.assertEqual(item[k].meta["patch_index"], i) + + def test_correct_center(self): + cropper = RandCropByPosNegLabeld(keys="label", label_key="label", spatial_size=[3, 3]) + cropper.set_random_state(0) + test_image = {"label": np.asarray([[[1, 0, 0, 1], [0, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 1]]])} + result = cropper(test_image) + np.testing.assert_allclose(result[0]["label"], np.asarray([[[0, 0, 1], [0, 0, 0], [0, 0, 0]]])) if __name__ == "__main__": diff --git a/tests/test_rand_cucim_dict_transform.py b/tests/test_rand_cucim_dict_transform.py index 2f27dd5f1f..a109bee845 100644 --- a/tests/test_rand_cucim_dict_transform.py +++ b/tests/test_rand_cucim_dict_transform.py @@ -91,12 +91,10 @@ class TestRandCuCIMDict(unittest.TestCase): ) def test_tramsforms_numpy_single(self, params, input, expected): input = {"image": np.copy(input)} - # apply_prob=1.0 output = RandCuCIMd(keys="image", apply_prob=1.0, **params)(input)["image"] self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, np.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIMd(keys="image", apply_prob=0.0, **params)(input)["image"] self.assertTrue(output.dtype == input["image"].dtype) self.assertTrue(isinstance(output, np.ndarray)) @@ -117,12 +115,10 @@ def test_tramsforms_numpy_single(self, params, input, expected): def test_tramsforms_numpy_batch(self, params, input, expected): input = {"image": np.copy(input[cp.newaxis, ...])} expected = expected[cp.newaxis, ...] - # apply_prob=1.0 output = RandCuCIMd(keys="image", apply_prob=1.0, **params)(input)["image"] self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, np.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIMd(keys="image", apply_prob=0.0, **params)(input)["image"] self.assertTrue(output.dtype == input["image"].dtype) self.assertTrue(isinstance(output, np.ndarray)) @@ -143,12 +139,10 @@ def test_tramsforms_numpy_batch(self, params, input, expected): def test_tramsforms_cupy_single(self, params, input, expected): input = {"image": cp.asarray(input)} expected = cp.asarray(expected) - # apply_prob=1.0 output = RandCuCIMd(keys="image", apply_prob=1.0, **params)(input)["image"] self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, cp.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIMd(keys="image", apply_prob=0.0, **params)(input)["image"] self.assertTrue(output.dtype == input["image"].dtype) self.assertTrue(isinstance(output, cp.ndarray)) @@ -169,12 +163,10 @@ def test_tramsforms_cupy_single(self, params, input, expected): def test_tramsforms_cupy_batch(self, params, input, expected): input = {"image": cp.asarray(input)[cp.newaxis, ...]} expected = cp.asarray(expected)[cp.newaxis, ...] - # apply_prob=1.0 output = RandCuCIMd(keys="image", **params)(input)["image"] self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, cp.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIMd(keys="image", apply_prob=0.0, **params)(input)["image"] self.assertTrue(output.dtype == input["image"].dtype) self.assertTrue(isinstance(output, cp.ndarray)) diff --git a/tests/test_rand_cucim_transform.py b/tests/test_rand_cucim_transform.py index c11ca5cd31..30164e4170 100644 --- a/tests/test_rand_cucim_transform.py +++ b/tests/test_rand_cucim_transform.py @@ -90,13 +90,11 @@ class TestRandCuCIM(unittest.TestCase): ] ) def test_tramsforms_numpy_single(self, params, input, expected): - # apply_prob=1.0 input = np.copy(input) output = RandCuCIM(apply_prob=1.0, **params)(input) self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, np.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIM(apply_prob=0.0, **params)(input) self.assertTrue(output.dtype == input.dtype) self.assertTrue(isinstance(output, np.ndarray)) @@ -117,12 +115,10 @@ def test_tramsforms_numpy_single(self, params, input, expected): def test_tramsforms_numpy_batch(self, params, input, expected): input = np.copy(input[cp.newaxis, ...]) expected = expected[cp.newaxis, ...] - # apply_prob=1.0 output = RandCuCIM(apply_prob=1.0, **params)(input) self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, np.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIM(apply_prob=0.0, **params)(input) self.assertTrue(output.dtype == input.dtype) self.assertTrue(isinstance(output, np.ndarray)) @@ -143,12 +139,10 @@ def test_tramsforms_numpy_batch(self, params, input, expected): def test_tramsforms_cupy_single(self, params, input, expected): input = cp.asarray(input) expected = cp.asarray(expected) - # apply_prob=1.0 output = RandCuCIM(apply_prob=1.0, **params)(input) self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, cp.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIM(apply_prob=0.0, **params)(input) self.assertTrue(output.dtype == input.dtype) self.assertTrue(isinstance(output, cp.ndarray)) @@ -169,12 +163,10 @@ def test_tramsforms_cupy_single(self, params, input, expected): def test_tramsforms_cupy_batch(self, params, input, expected): input = cp.asarray(input)[cp.newaxis, ...] expected = cp.asarray(expected)[cp.newaxis, ...] - # apply_prob=1.0 output = RandCuCIM(**params)(input) self.assertTrue(output.dtype == expected.dtype) self.assertTrue(isinstance(output, cp.ndarray)) cp.testing.assert_allclose(output, expected) - # apply_prob=0.0 output = RandCuCIM(apply_prob=0.0, **params)(input) self.assertTrue(output.dtype == input.dtype) self.assertTrue(isinstance(output, cp.ndarray)) diff --git a/tests/test_rand_deform_grid.py b/tests/test_rand_deform_grid.py index 8a2c8bf6eb..3e59e3207b 100644 --- a/tests/test_rand_deform_grid.py +++ b/tests/test_rand_deform_grid.py @@ -19,7 +19,7 @@ TEST_CASES = [ [ - dict(spacing=(1, 2), magnitude_range=(1.0, 2.0), as_tensor_output=False, device=None), + dict(spacing=(1, 2), magnitude_range=(1.0, 2.0), device=None), {"spatial_size": (3, 3)}, np.array( [ @@ -48,7 +48,7 @@ ), ], [ - dict(spacing=(1, 2, 2), magnitude_range=(1.0, 3.0), as_tensor_output=False, device=None), + dict(spacing=(1, 2, 2), magnitude_range=(1.0, 3.0), device=None), {"spatial_size": (1, 2, 2)}, np.array( [ diff --git a/tests/test_rand_elastic_2d.py b/tests/test_rand_elastic_2d.py index bc23a6c5cb..125da74528 100644 --- a/tests/test_rand_elastic_2d.py +++ b/tests/test_rand_elastic_2d.py @@ -15,13 +15,14 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Rand2DElastic -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, is_tf32_env _rtol = 5e-3 if is_tf32_env() else 1e-4 TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -110,9 +111,14 @@ class TestRand2DElastic(unittest.TestCase): @parameterized.expand(TESTS) def test_rand_2d_elastic(self, input_param, input_data, expected_val): g = Rand2DElastic(**input_param) + set_track_meta(False) + result = g(**input_data) + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) g.set_random_state(123) result = g(**input_data) - assert_allclose(result, expected_val, rtol=_rtol, atol=1e-4) + assert_allclose(result, expected_val, type_test=False, rtol=_rtol, atol=1e-4) if __name__ == "__main__": diff --git a/tests/test_rand_elastic_3d.py b/tests/test_rand_elastic_3d.py index 39ce779cb0..76c9e9024d 100644 --- a/tests/test_rand_elastic_3d.py +++ b/tests/test_rand_elastic_3d.py @@ -15,11 +15,12 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Rand3DElastic -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -86,9 +87,15 @@ class TestRand3DElastic(unittest.TestCase): @parameterized.expand(TESTS) def test_rand_3d_elastic(self, input_param, input_data, expected_val): g = Rand3DElastic(**input_param) + set_track_meta(False) g.set_random_state(123) result = g(**input_data) - assert_allclose(result, expected_val, rtol=1e-1, atol=1e-1) + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) + g.set_random_state(123) + result = g(**input_data) + assert_allclose(result, expected_val, type_test=False, rtol=1e-1, atol=1e-1) if __name__ == "__main__": diff --git a/tests/test_rand_elasticd_2d.py b/tests/test_rand_elasticd_2d.py index ead39e5731..759ba2c4da 100644 --- a/tests/test_rand_elasticd_2d.py +++ b/tests/test_rand_elasticd_2d.py @@ -16,12 +16,12 @@ from parameterized import parameterized from monai.transforms import Rand2DElasticd -from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose, is_tf32_env _rtol = 5e-3 if is_tf32_env() else 1e-4 TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -166,7 +166,7 @@ def test_rand_2d_elasticd(self, input_param, input_data, expected_val): for key in res: result = res[key] expected = expected_val[key] if isinstance(expected_val, dict) else expected_val - assert_allclose(result, expected, rtol=_rtol, atol=5e-3) + assert_allclose(result, expected, rtol=_rtol, atol=5e-3, type_test=False) if __name__ == "__main__": diff --git a/tests/test_rand_elasticd_3d.py b/tests/test_rand_elasticd_3d.py index c78ed1f42e..eaba06c953 100644 --- a/tests/test_rand_elasticd_3d.py +++ b/tests/test_rand_elasticd_3d.py @@ -16,10 +16,10 @@ from parameterized import parameterized from monai.transforms import Rand3DElasticd -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -145,7 +145,7 @@ def test_rand_3d_elasticd(self, input_param, input_data, expected_val): for key in res: result = res[key] expected = expected_val[key] if isinstance(expected_val, dict) else expected_val - assert_allclose(result, expected, rtol=1e-2, atol=1e-2) + assert_allclose(result, expected, type_test=False, rtol=1e-2, atol=1e-2) if __name__ == "__main__": diff --git a/tests/test_rand_flip.py b/tests/test_rand_flip.py index b9e9a8c4d6..cdd51dd77e 100644 --- a/tests/test_rand_flip.py +++ b/tests/test_rand_flip.py @@ -12,10 +12,12 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandFlip -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion INVALID_CASES = [("wrong_axis", ["s", 1], TypeError), ("not_numbers", "s", TypeError)] @@ -31,13 +33,19 @@ def test_invalid_inputs(self, _, spatial_axis, raises): @parameterized.expand(VALID_CASES) def test_correct_results(self, _, spatial_axis): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: im = p(self.imt[0]) flip = RandFlip(prob=1.0, spatial_axis=spatial_axis) + set_track_meta(False) + result = flip(im) + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) expected = [np.flip(channel, spatial_axis) for channel in self.imt[0]] expected = np.stack(expected) result = flip(im) - assert_allclose(result, p(expected)) + assert_allclose(result, p(expected), type_test="tensor") + test_local_inversion(flip, result, im) if __name__ == "__main__": diff --git a/tests/test_rand_flipd.py b/tests/test_rand_flipd.py index 9a92661c59..92b070fd0a 100644 --- a/tests/test_rand_flipd.py +++ b/tests/test_rand_flipd.py @@ -12,10 +12,12 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandFlipd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion VALID_CASES = [("no_axis", None), ("one_axis", 1), ("many_axis", [0, 1])] @@ -23,12 +25,19 @@ class TestRandFlipd(NumpyImageTestCase2D): @parameterized.expand(VALID_CASES) def test_correct_results(self, _, spatial_axis): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: flip = RandFlipd(keys="img", prob=1.0, spatial_axis=spatial_axis) - result = flip({"img": p(self.imt[0])})["img"] + im = p(self.imt[0]) + result = flip({"img": im})["img"] expected = [np.flip(channel, spatial_axis) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(result, p(expected)) + assert_allclose(result, p(expected), type_test="tensor") + test_local_inversion(flip, {"img": result}, {"img": im}, "img") + set_track_meta(False) + result = flip({"img": im})["img"] + self.assertNotIsInstance(result, MetaTensor) + self.assertIsInstance(result, torch.Tensor) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_rand_gaussian_noise.py b/tests/test_rand_gaussian_noise.py index 1f2adfb9e7..faa2b143f5 100644 --- a/tests/test_rand_gaussian_noise.py +++ b/tests/test_rand_gaussian_noise.py @@ -35,7 +35,6 @@ def test_correct_results(self, _, im_type, mean, std): np.random.seed(seed) np.random.random() expected = self.imt + np.random.normal(mean, np.random.uniform(0, std), size=self.imt.shape) - self.assertEqual(type(im), type(noised)) if isinstance(noised, torch.Tensor): noised = noised.cpu() np.testing.assert_allclose(expected, noised, atol=1e-5) diff --git a/tests/test_rand_gaussian_noised.py b/tests/test_rand_gaussian_noised.py index be1df0f2e6..a927761186 100644 --- a/tests/test_rand_gaussian_noised.py +++ b/tests/test_rand_gaussian_noised.py @@ -39,7 +39,6 @@ def test_correct_results(self, _, im_type, keys, mean, std): noise = np.random.normal(mean, np.random.uniform(0, std), size=self.imt.shape) for k in keys: expected = self.imt + noise - self.assertEqual(type(im), type(noised[k])) if isinstance(noised[k], torch.Tensor): noised[k] = noised[k].cpu() np.testing.assert_allclose(expected, noised[k], atol=1e-5, rtol=1e-5) diff --git a/tests/test_rand_gaussian_sharpen.py b/tests/test_rand_gaussian_sharpen.py index 06563a35b6..3f7f276cd3 100644 --- a/tests/test_rand_gaussian_sharpen.py +++ b/tests/test_rand_gaussian_sharpen.py @@ -131,7 +131,7 @@ def test_value(self, argments, image, expected_data): converter = RandGaussianSharpen(**argments) converter.set_random_state(seed=0) result = converter(image) - assert_allclose(result, expected_data, atol=0, rtol=1e-4, type_test=False) + assert_allclose(result, expected_data, atol=0, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_gaussian_smooth.py b/tests/test_rand_gaussian_smooth.py index d51618be95..e395d4f395 100644 --- a/tests/test_rand_gaussian_smooth.py +++ b/tests/test_rand_gaussian_smooth.py @@ -89,7 +89,7 @@ def test_value(self, argments, image, expected_data): converter = RandGaussianSmooth(**argments) converter.set_random_state(seed=0) result = converter(image) - assert_allclose(result, expected_data, rtol=1e-4, type_test=False) + assert_allclose(result, expected_data, rtol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_gibbs_noise.py b/tests/test_rand_gibbs_noise.py index fe928038da..b87b839eb9 100644 --- a/tests/test_rand_gibbs_noise.py +++ b/tests/test_rand_gibbs_noise.py @@ -13,14 +13,13 @@ from copy import deepcopy import numpy as np -import torch from parameterized import parameterized from monai.data.synthetic import create_test_image_2d, create_test_image_3d from monai.transforms import RandGibbsNoise from monai.utils.misc import set_determinism from monai.utils.module import optional_import -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose _, has_torch_fft = optional_import("torch.fft", name="fftshift") @@ -50,7 +49,7 @@ def test_0_prob(self, im_shape, input_type): alpha = [0.5, 1.0] t = RandGibbsNoise(0.0, alpha) out = t(im) - torch.testing.assert_allclose(im, out, rtol=1e-7, atol=0) + assert_allclose(out, im, rtol=1e-7, atol=0, type_test="tensor") @parameterized.expand(TEST_CASES) def test_same_result(self, im_shape, input_type): @@ -61,8 +60,7 @@ def test_same_result(self, im_shape, input_type): out1 = t(deepcopy(im)) t.set_random_state(42) out2 = t(deepcopy(im)) - torch.testing.assert_allclose(out1, out2, rtol=1e-7, atol=0) - self.assertIsInstance(out1, type(im)) + assert_allclose(out1, out2, rtol=1e-7, atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_identity(self, im_shape, input_type): @@ -70,7 +68,7 @@ def test_identity(self, im_shape, input_type): alpha = [0.0, 0.0] t = RandGibbsNoise(1.0, alpha) out = t(deepcopy(im)) - torch.testing.assert_allclose(im, out, atol=1e-2, rtol=1e-7) + assert_allclose(out, im, atol=1e-2, rtol=1e-7, type_test="tensor") @parameterized.expand(TEST_CASES) def test_alpha_1(self, im_shape, input_type): @@ -78,7 +76,7 @@ def test_alpha_1(self, im_shape, input_type): alpha = [1.0, 1.0] t = RandGibbsNoise(1.0, alpha) out = t(deepcopy(im)) - torch.testing.assert_allclose(0 * im, out, rtol=1e-7, atol=0) + assert_allclose(out, 0 * im, rtol=1e-7, atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_alpha(self, im_shape, input_type): diff --git a/tests/test_rand_gibbs_noised.py b/tests/test_rand_gibbs_noised.py index 8c5e045b90..8b15fcc267 100644 --- a/tests/test_rand_gibbs_noised.py +++ b/tests/test_rand_gibbs_noised.py @@ -20,7 +20,7 @@ from monai.transforms import RandGibbsNoised from monai.utils.misc import set_determinism from monai.utils.module import optional_import -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose _, has_torch_fft = optional_import("torch.fft", name="fftshift") @@ -65,8 +65,7 @@ def test_same_result(self, im_shape, input_type): t.set_random_state(42) out2 = t(deepcopy(data)) for k in KEYS: - torch.testing.assert_allclose(out1[k], out2[k], rtol=1e-7, atol=0) - self.assertIsInstance(out1[k], type(data[k])) + assert_allclose(out1[k], out2[k], rtol=1e-7, atol=0, type_test="tensor") @parameterized.expand(TEST_CASES) def test_identity(self, im_shape, input_type): @@ -75,11 +74,7 @@ def test_identity(self, im_shape, input_type): t = RandGibbsNoised(KEYS, 1.0, alpha) out = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) - if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) - out[k], data[k] = out[k].cpu(), data[k].cpu() - np.testing.assert_allclose(data[k], out[k], atol=1e-2) + assert_allclose(out[k], data[k], atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_alpha_1(self, im_shape, input_type): @@ -88,11 +83,7 @@ def test_alpha_1(self, im_shape, input_type): t = RandGibbsNoised(KEYS, 1.0, alpha) out = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out[k]), type(data[k])) - if isinstance(out[k], torch.Tensor): - self.assertEqual(out[k].device, data[k].device) - out[k], data[k] = out[k].cpu(), data[k].cpu() - np.testing.assert_allclose(0.0 * data[k], out[k], atol=1e-2) + assert_allclose(out[k], 0.0 * data[k], atol=1e-2, type_test="tensor") @parameterized.expand(TEST_CASES) def test_dict_matches(self, im_shape, input_type): diff --git a/tests/test_rand_grid_distortion.py b/tests/test_rand_grid_distortion.py index 80f19df0db..88b4989cd5 100644 --- a/tests/test_rand_grid_distortion.py +++ b/tests/test_rand_grid_distortion.py @@ -15,10 +15,10 @@ from parameterized import parameterized from monai.transforms import RandGridDistortion -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: seed = 0 TESTS.append( [ @@ -87,7 +87,7 @@ def test_rand_grid_distortion(self, input_param, seed, input_data, expected_val) g = RandGridDistortion(**input_param) g.set_random_state(seed=seed) result = g(input_data) - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4) + assert_allclose(result, expected_val, type_test="tensor", rtol=1e-4, atol=1e-4) if __name__ == "__main__": diff --git a/tests/test_rand_grid_distortiond.py b/tests/test_rand_grid_distortiond.py index 323848dc0b..a7b64e5980 100644 --- a/tests/test_rand_grid_distortiond.py +++ b/tests/test_rand_grid_distortiond.py @@ -15,12 +15,12 @@ from parameterized import parameterized from monai.transforms import RandGridDistortiond -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] num_cells = 2 seed = 0 -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: img = np.indices([6, 6]).astype(np.float32) TESTS.append( [ @@ -80,8 +80,8 @@ def test_rand_grid_distortiond(self, input_param, seed, input_data, expected_val g = RandGridDistortiond(**input_param) g.set_random_state(seed=seed) result = g(input_data) - assert_allclose(result["img"], expected_val_img, rtol=1e-4, atol=1e-4) - assert_allclose(result["mask"], expected_val_mask, rtol=1e-4, atol=1e-4) + assert_allclose(result["img"], expected_val_img, type_test=False, rtol=1e-4, atol=1e-4) + assert_allclose(result["mask"], expected_val_mask, type_test=False, rtol=1e-4, atol=1e-4) if __name__ == "__main__": diff --git a/tests/test_rand_grid_patch.py b/tests/test_rand_grid_patch.py new file mode 100644 index 0000000000..3957dc1ce8 --- /dev/null +++ b/tests/test_rand_grid_patch.py @@ -0,0 +1,88 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.spatial.array import RandGridPatch +from monai.utils import set_determinism +from tests.utils import TEST_NDARRAYS, assert_allclose + +set_determinism(1234) + +A = np.arange(16).repeat(3).reshape(4, 4, 3).transpose(2, 0, 1) +A11 = A[:, :2, :2] +A12 = A[:, :2, 2:] +A21 = A[:, 2:, :2] +A22 = A[:, 2:, 2:] + +TEST_CASE_0 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0}, A, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"patch_size": (2, 2), "min_offset": 0, "num_patches": 3}, A, [A11, A12, A21]] +TEST_CASE_2 = [ + {"patch_size": (2, 2), "min_offset": 0, "max_offset": 0, "num_patches": 5}, + A, + [A11, A12, A21, A22, np.zeros((3, 2, 2))], +] +TEST_CASE_3 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0}, A, [A11, A12, A21, A22]] +TEST_CASE_4 = [{"patch_size": (2, 2)}, A, [A11, A12, A21, A22]] +TEST_CASE_5 = [{"patch_size": (2, 2), "min_offset": 2, "max_offset": 2}, A, [A22]] +TEST_CASE_6 = [{"patch_size": (2, 2), "min_offset": (0, 2), "max_offset": (0, 2)}, A, [A12, A22]] +TEST_CASE_7 = [{"patch_size": (2, 2), "min_offset": 1, "max_offset": 2}, A, [A22]] +TEST_CASE_8 = [ + {"patch_size": (2, 2), "min_offset": 0, "max_offset": 1, "num_patches": 1, "sort_fn": "max"}, + A, + [A[:, 1:3, 1:3]], +] +TEST_CASE_9 = [ + { + "patch_size": (3, 3), + "min_offset": -3, + "max_offset": -1, + "sort_fn": "min", + "num_patches": 1, + "constant_values": 255, + }, + A, + [np.pad(A[:, :2, 1:], ((0, 0), (1, 0), (0, 0)), mode="constant", constant_values=255)], +] +TEST_CASE_10 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0, "threshold": 50.0}, A, [A11]] + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + TEST_SINGLE.append([p, *TEST_CASE_9]) + TEST_SINGLE.append([p, *TEST_CASE_10]) + + +class TestRandGridPatch(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_rand_grid_patch(self, in_type, input_parameters, image, expected): + input_image = in_type(image) + splitter = RandGridPatch(**input_parameters) + splitter.set_random_state(1234) + output = list(splitter(input_image)) + self.assertEqual(len(output), len(expected)) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[0], expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_rand_grid_patchd.py b/tests/test_rand_grid_patchd.py new file mode 100644 index 0000000000..656fbd9e36 --- /dev/null +++ b/tests/test_rand_grid_patchd.py @@ -0,0 +1,93 @@ +# 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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.spatial.dictionary import RandGridPatchd +from monai.utils import set_determinism +from tests.utils import TEST_NDARRAYS, assert_allclose + +set_determinism(1234) + +A = np.arange(16).repeat(3).reshape(4, 4, 3).transpose(2, 0, 1) +A11 = A[:, :2, :2] +A12 = A[:, :2, 2:] +A21 = A[:, 2:, :2] +A22 = A[:, 2:, 2:] + +TEST_CASE_0 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_1 = [{"patch_size": (2, 2), "min_offset": 0, "num_patches": 3}, {"image": A}, [A11, A12, A21]] +TEST_CASE_2 = [ + {"patch_size": (2, 2), "min_offset": 0, "max_offset": 0, "num_patches": 5}, + {"image": A}, + [A11, A12, A21, A22, np.zeros((3, 2, 2))], +] +TEST_CASE_3 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_4 = [{"patch_size": (2, 2)}, {"image": A}, [A11, A12, A21, A22]] +TEST_CASE_5 = [{"patch_size": (2, 2), "min_offset": 2, "max_offset": 2}, {"image": A}, [A22]] +TEST_CASE_6 = [{"patch_size": (2, 2), "min_offset": (0, 2), "max_offset": (0, 2)}, {"image": A}, [A12, A22]] +TEST_CASE_7 = [{"patch_size": (2, 2), "min_offset": 1, "max_offset": 2}, {"image": A}, [A22]] +TEST_CASE_8 = [ + {"patch_size": (2, 2), "min_offset": 0, "max_offset": 1, "num_patches": 1, "sort_fn": "max"}, + {"image": A}, + [A[:, 1:3, 1:3]], +] +TEST_CASE_9 = [ + { + "patch_size": (3, 3), + "min_offset": -3, + "max_offset": -1, + "sort_fn": "min", + "num_patches": 1, + "constant_values": 255, + }, + {"image": A}, + [np.pad(A[:, :2, 1:], ((0, 0), (1, 0), (0, 0)), mode="constant", constant_values=255)], +] +TEST_CASE_10 = [{"patch_size": (2, 2), "min_offset": 0, "max_offset": 0, "threshold": 50.0}, {"image": A}, [A11]] + +TEST_SINGLE = [] +for p in TEST_NDARRAYS: + TEST_SINGLE.append([p, *TEST_CASE_0]) + TEST_SINGLE.append([p, *TEST_CASE_1]) + TEST_SINGLE.append([p, *TEST_CASE_2]) + TEST_SINGLE.append([p, *TEST_CASE_3]) + TEST_SINGLE.append([p, *TEST_CASE_4]) + TEST_SINGLE.append([p, *TEST_CASE_5]) + TEST_SINGLE.append([p, *TEST_CASE_6]) + TEST_SINGLE.append([p, *TEST_CASE_7]) + TEST_SINGLE.append([p, *TEST_CASE_8]) + TEST_SINGLE.append([p, *TEST_CASE_9]) + TEST_SINGLE.append([p, *TEST_CASE_10]) + + +class TestRandGridPatchd(unittest.TestCase): + @parameterized.expand(TEST_SINGLE) + def test_rand_grid_patchd(self, in_type, input_parameters, image_dict, expected): + image_key = "image" + input_dict = {} + for k, v in image_dict.items(): + input_dict[k] = v + if k == image_key: + input_dict[k] = in_type(v) + splitter = RandGridPatchd(keys=image_key, **input_parameters) + splitter.set_random_state(1234) + output = list(splitter(input_dict)) + self.assertEqual(len(output), len(expected)) + for output_patch, expected_patch in zip(output, expected): + assert_allclose(output_patch[image_key], expected_patch, type_test=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_rand_histogram_shift.py b/tests/test_rand_histogram_shift.py index c66f7859c6..89198549cd 100644 --- a/tests/test_rand_histogram_shift.py +++ b/tests/test_rand_histogram_shift.py @@ -12,6 +12,7 @@ import unittest import numpy as np +import torch from parameterized import parameterized from monai.transforms import RandHistogramShift @@ -48,7 +49,25 @@ def test_rand_histogram_shift(self, input_param, input_data, expected_val): g = RandHistogramShift(**input_param) g.set_random_state(123) result = g(**input_data) - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4, type_test=False) + assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4, type_test="tensor") + + def test_interp(self): + tr = RandHistogramShift() + for array_type in (torch.tensor, np.array): + x = array_type([0.0, 4.0, 6.0, 10.0]) + y = array_type([1.0, -1.0, 3.0, 5.0]) + + yi = tr.interp(array_type([0, 2, 4, 8, 10]), x, y) + self.assertEqual(yi.shape, (5,)) + assert_allclose(yi, array_type([1.0, 0.0, -1.0, 4.0, 5.0])) + + yi = tr.interp(array_type([-1, 11, 10.001, -0.001]), x, y) + self.assertEqual(yi.shape, (4,)) + assert_allclose(yi, array_type([1.0, 5.0, 5.0, 1.0])) + + yi = tr.interp(array_type([[-2, 11], [1, 3], [8, 10]]), x, y) + self.assertEqual(yi.shape, (3, 2)) + assert_allclose(yi, array_type([[1.0, 5.0], [0.5, -0.5], [4.0, 5.0]])) if __name__ == "__main__": diff --git a/tests/test_rand_histogram_shiftd.py b/tests/test_rand_histogram_shiftd.py index fe8ddf9ffd..7c94379e0e 100644 --- a/tests/test_rand_histogram_shiftd.py +++ b/tests/test_rand_histogram_shiftd.py @@ -64,10 +64,10 @@ def test_rand_histogram_shiftd(self, input_param, input_data, expected_val): g = RandHistogramShiftd(**input_param) g.set_random_state(123) res = g(input_data) - for key in res: + for key in ("img",): result = res[key] expected = expected_val[key] if isinstance(expected_val, dict) else expected_val - assert_allclose(result, expected, rtol=1e-4, atol=1e-4, type_test=False) + assert_allclose(result, expected, rtol=1e-4, atol=1e-4, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_k_space_spike_noise.py b/tests/test_rand_k_space_spike_noise.py index 8027194555..176699ddd1 100644 --- a/tests/test_rand_k_space_spike_noise.py +++ b/tests/test_rand_k_space_spike_noise.py @@ -12,14 +12,12 @@ import unittest from copy import deepcopy -import numpy as np -import torch from parameterized import parameterized from monai.data.synthetic import create_test_image_2d, create_test_image_3d from monai.transforms import KSpaceSpikeNoise, RandKSpaceSpikeNoise from monai.utils.misc import set_determinism -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose TESTS = [] for shape in ((128, 64), (64, 48, 80)): @@ -48,11 +46,7 @@ def test_0_prob(self, im_shape, im_type, channel_wise): intensity_range = [14, 15] t = RandKSpaceSpikeNoise(0.0, intensity_range, channel_wise) out = t(im) - self.assertEqual(type(im), type(out)) - if isinstance(out, torch.Tensor): - self.assertEqual(out.device, im.device) - im, out = im.cpu(), out.cpu() - np.testing.assert_allclose(im, out) + assert_allclose(out, im, type_test="tensor") @parameterized.expand(TESTS) def test_1_prob(self, im_shape, im_type, channel_wise): @@ -62,11 +56,7 @@ def test_1_prob(self, im_shape, im_type, channel_wise): out = t(im) base_t = KSpaceSpikeNoise(t.sampled_locs, [14]) out = out - base_t(im) - self.assertEqual(type(im), type(out)) - if isinstance(out, torch.Tensor): - self.assertEqual(out.device, im.device) - im, out = im.cpu(), out.cpu() - np.testing.assert_allclose(out, im * 0) + assert_allclose(out, im * 0, type_test="tensor") @parameterized.expand(TESTS) def test_same_result(self, im_shape, im_type, channel_wise): @@ -77,11 +67,7 @@ def test_same_result(self, im_shape, im_type, channel_wise): out1 = t(deepcopy(im)) t.set_random_state(42) out2 = t(deepcopy(im)) - self.assertEqual(type(im), type(out1)) - if isinstance(out1, torch.Tensor): - self.assertEqual(out1.device, im.device) - out1, out2 = out1.cpu(), out2.cpu() - np.testing.assert_allclose(out1, out2) + assert_allclose(out1, out2, type_test="tensor") @parameterized.expand(TESTS) def test_intensity(self, im_shape, im_type, channel_wise): diff --git a/tests/test_rand_k_space_spike_noised.py b/tests/test_rand_k_space_spike_noised.py index 7a6a73b215..156c95822f 100644 --- a/tests/test_rand_k_space_spike_noised.py +++ b/tests/test_rand_k_space_spike_noised.py @@ -12,14 +12,12 @@ import unittest from copy import deepcopy -import numpy as np -import torch from parameterized import parameterized from monai.data.synthetic import create_test_image_2d, create_test_image_3d from monai.transforms import RandKSpaceSpikeNoised from monai.utils.misc import set_determinism -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose TESTS = [] for shape in ((128, 64), (64, 48, 80)): @@ -57,33 +55,20 @@ def test_same_result(self, im_shape, im_type): out2 = t(deepcopy(data)) for k in KEYS: - self.assertEqual(type(out1[k]), type(data[k])) - if isinstance(out1[k], torch.Tensor): - self.assertEqual(out1[k].device, data[k].device) - out1[k] = out1[k].cpu() - out2[k] = out2[k].cpu() - np.testing.assert_allclose(out1[k], out2[k], atol=1e-10) + assert_allclose(out1[k], out2[k], atol=1e-10, type_test="tensor") @parameterized.expand(TESTS) def test_0_prob(self, im_shape, im_type): data = self.get_data(im_shape, im_type) t1 = RandKSpaceSpikeNoised(KEYS, prob=0.0, intensity_range=(13, 15), channel_wise=True) - t2 = RandKSpaceSpikeNoised(KEYS, prob=0.0, intensity_range=(13, 15), channel_wise=True) out1 = t1(data) out2 = t2(data) for k in KEYS: - self.assertEqual(type(out1[k]), type(data[k])) - if isinstance(out1[k], torch.Tensor): - self.assertEqual(out1[k].device, data[k].device) - out1[k] = out1[k].cpu() - out2[k] = out2[k].cpu() - data[k] = data[k].cpu() - - np.testing.assert_allclose(data[k], out1[k]) - np.testing.assert_allclose(data[k], out2[k]) + assert_allclose(out1[k], data[k], type_test="tensor") + assert_allclose(out2[k], data[k], type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_lambda.py b/tests/test_rand_lambda.py index 043f44aec4..c356406f61 100644 --- a/tests/test_rand_lambda.py +++ b/tests/test_rand_lambda.py @@ -10,11 +10,15 @@ # limitations under the License. import unittest +from copy import deepcopy import numpy as np +from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms.transform import Randomizable from monai.transforms.utility.array import RandLambda +from tests.utils import TEST_NDARRAYS, assert_allclose class RandTest(Randomizable): @@ -27,26 +31,56 @@ def randomize(self, data=None): def __call__(self, data): self.randomize() - return data + self._a + return deepcopy(data) + self._a class TestRandLambda(unittest.TestCase): - def test_rand_lambdad_identity(self): - img = np.zeros((10, 10)) + def check(self, tr: RandLambda, img, img_orig_type, out, expected=None): + # input shouldn't change + self.assertIsInstance(img, img_orig_type) + if isinstance(img, MetaTensor): + self.assertEqual(len(img.applied_operations), 0) + # output data matches expected + assert_allclose(expected, out, type_test=False) + # output type is MetaTensor with 1 appended operation + self.assertIsInstance(out, MetaTensor) + self.assertEqual(len(out.applied_operations), 1) + + # inverse + inv = tr.inverse(out) + # after inverse, input image remains unchanged + self.assertIsInstance(img, img_orig_type) + if isinstance(img, MetaTensor): + self.assertEqual(len(img.applied_operations), 0) + # after inverse, output is MetaTensor with 0 applied operations + self.assertIsInstance(inv, MetaTensor) + self.assertEqual(len(inv.applied_operations), 0) + + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_rand_lambdad_identity(self, t): + img = t(np.zeros((10, 10))) + img_t = type(img) test_func = RandTest() test_func.set_random_state(seed=134) expected = test_func(img) test_func.set_random_state(seed=134) - ret = RandLambda(func=test_func)(img) - np.testing.assert_allclose(expected, ret) - ret = RandLambda(func=test_func, prob=0.0)(img) - np.testing.assert_allclose(img, ret) + # default prob + tr = RandLambda(func=test_func) + ret = tr(img) + self.check(tr, img, img_t, ret, expected) + + # prob = 0 + tr = RandLambda(func=test_func, prob=0.0) + ret = tr(img) + self.check(tr, img, img_t, ret, expected=img) + + # prob = 0.5 trans = RandLambda(func=test_func, prob=0.5) trans.set_random_state(seed=123) ret = trans(img) - np.testing.assert_allclose(img, ret) + self.check(trans, img, img_t, ret, expected=img) if __name__ == "__main__": diff --git a/tests/test_rand_lambdad.py b/tests/test_rand_lambdad.py index 854fef8879..b181db5035 100644 --- a/tests/test_rand_lambdad.py +++ b/tests/test_rand_lambdad.py @@ -10,11 +10,15 @@ # limitations under the License. import unittest +from copy import deepcopy import numpy as np +from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms.transform import Randomizable from monai.transforms.utility.dictionary import RandLambdad +from tests.utils import TEST_NDARRAYS, assert_allclose class RandTest(Randomizable): @@ -31,26 +35,42 @@ def __call__(self, data): class TestRandLambdad(unittest.TestCase): - def test_rand_lambdad_identity(self): - img = np.zeros((10, 10)) + def check(self, tr: RandLambdad, input: dict, out: dict, expected: dict): + if isinstance(input["img"], MetaTensor): + self.assertEqual(len(input["img"].applied_operations), 0) + self.assertIsInstance(out["img"], MetaTensor) + self.assertEqual(len(out["img"].applied_operations), 1) + assert_allclose(expected["img"], out["img"], type_test=False) + assert_allclose(expected["prop"], out["prop"], type_test=False) + inv = tr.inverse(out) + self.assertIsInstance(inv["img"], MetaTensor) + self.assertEqual(len(inv["img"].applied_operations), 0) # type: ignore + + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_rand_lambdad_identity(self, t): + img = t(np.zeros((10, 10))) data = {"img": img, "prop": 1.0} test_func = RandTest() test_func.set_random_state(seed=134) expected = {"img": test_func(data["img"]), "prop": 1.0} test_func.set_random_state(seed=134) - ret = RandLambdad(keys=["img", "prop"], func=test_func, overwrite=[True, False])(data) - np.testing.assert_allclose(expected["img"], ret["img"]) - np.testing.assert_allclose(expected["prop"], ret["prop"]) - ret = RandLambdad(keys=["img", "prop"], func=test_func, prob=0.0)(data) - np.testing.assert_allclose(data["img"], ret["img"]) - np.testing.assert_allclose(data["prop"], ret["prop"]) + # default prob + tr = RandLambdad(keys=["img", "prop"], func=test_func, overwrite=[True, False]) + ret = tr(deepcopy(data)) + self.check(tr, data, ret, expected) + + # prob = 0 + tr = RandLambdad(keys=["img", "prop"], func=test_func, prob=0.0) + ret = tr(deepcopy(data)) + self.check(tr, data, ret, expected=data) + + # prob = 0.5 trans = RandLambdad(keys=["img", "prop"], func=test_func, prob=0.5) trans.set_random_state(seed=123) - ret = trans(data) - np.testing.assert_allclose(data["img"], ret["img"]) - np.testing.assert_allclose(data["prop"], ret["prop"]) + ret = trans(deepcopy(data)) + self.check(trans, data, ret, expected=data) if __name__ == "__main__": diff --git a/tests/test_rand_rician_noise.py b/tests/test_rand_rician_noise.py index 8e2ea1ee3a..896ae8b2e0 100644 --- a/tests/test_rand_rician_noise.py +++ b/tests/test_rand_rician_noise.py @@ -30,7 +30,8 @@ def test_correct_results(self, _, in_type, mean, std): seed = 0 rician_fn = RandRicianNoise(prob=1.0, mean=mean, std=std) rician_fn.set_random_state(seed) - noised = rician_fn(in_type(self.imt)) + im = in_type(self.imt) + noised = rician_fn(im) np.random.seed(seed) np.random.random() _std = np.random.uniform(0, std) diff --git a/tests/test_rand_rotate.py b/tests/test_rand_rotate.py index 7a85fce23b..bdee0474d0 100644 --- a/tests/test_rand_rotate.py +++ b/tests/test_rand_rotate.py @@ -17,18 +17,19 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandRotate -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D, test_local_inversion TEST_CASES_2D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_2D.append((p, np.pi / 2, True, "bilinear", "border", False)) TEST_CASES_2D.append((p, np.pi / 4, True, "nearest", "border", False)) TEST_CASES_2D.append((p, np.pi, False, "nearest", "zeros", True)) TEST_CASES_2D.append((p, (-np.pi / 4, 0), False, "nearest", "zeros", True)) TEST_CASES_3D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_3D.append( (p, np.pi / 2, -np.pi / 6, (0.0, np.pi), False, "bilinear", "border", False, (1, 87, 104, 109)) ) @@ -108,8 +109,16 @@ def test_correct_results(self, im_type, x, y, z, keep_size, mode, padding_mode, dtype=np.float64, ) rotate_fn.set_random_state(243) - rotated = rotate_fn(im_type(self.imt[0])) + im = im_type(self.imt[0]) + rotated = rotate_fn(im) torch.testing.assert_allclose(rotated.shape, expected, rtol=1e-7, atol=0) + test_local_inversion(rotate_fn, rotated, im) + + set_track_meta(False) + rotated = rotate_fn(im) + self.assertNotIsInstance(rotated, MetaTensor) + self.assertIsInstance(rotated, torch.Tensor) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_rand_rotate90.py b/tests/test_rand_rotate90.py index b845944062..30ad906ac2 100644 --- a/tests/test_rand_rotate90.py +++ b/tests/test_rand_rotate90.py @@ -12,47 +12,64 @@ import unittest import numpy as np +import torch +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandRotate90 -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion class TestRandRotate90(NumpyImageTestCase2D): def test_default(self): rotate = RandRotate90() - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate.set_random_state(123) - rotated = rotate(p(self.imt[0])) + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") def test_k(self): rotate = RandRotate90(max_k=2) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + set_track_meta(False) + rotated = rotate(im) + self.assertNotIsInstance(rotated, MetaTensor) + self.assertIsInstance(rotated, torch.Tensor) + + set_track_meta(True) rotate.set_random_state(123) - rotated = rotate(p(self.imt[0])) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") def test_spatial_axes(self): - rotate = RandRotate90(spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: - rotate.set_random_state(123) - rotated = rotate(p(self.imt[0])) - expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] + rotate = RandRotate90(spatial_axes=(0, 1), prob=1.0) + for p in TEST_NDARRAYS_ALL: + rotate.set_random_state(1234) + im = p(self.imt[0]) + rotated = rotate(im) + self.assertEqual(len(rotated.applied_operations), 1) + expected = [np.rot90(channel, rotate._rand_k, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + test_local_inversion(rotate, rotated, im) def test_prob_k_spatial_axes(self): rotate = RandRotate90(prob=1.0, max_k=2, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate.set_random_state(234) - rotated = rotate(p(self.imt[0])) + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_rotate90d.py b/tests/test_rand_rotate90d.py index ded18e430a..ec0e5ac92e 100644 --- a/tests/test_rand_rotate90d.py +++ b/tests/test_rand_rotate90d.py @@ -12,51 +12,67 @@ import unittest import numpy as np +import torch +from monai.data import MetaTensor, set_track_meta from monai.transforms import RandRotate90d -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion class TestRandRotate90d(NumpyImageTestCase2D): def test_default(self): key = None rotate = RandRotate90d(keys=key) - for p in TEST_NDARRAYS: - rotate.set_random_state(123) - rotated = rotate({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + rotate.set_random_state(1323) + im = {key: p(self.imt[0])} + rotated = rotate(im) + test_local_inversion(rotate, rotated, im, key) expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") + + set_track_meta(False) + rotated = rotate(im)[key] + self.assertNotIsInstance(rotated, MetaTensor) + self.assertIsInstance(rotated, torch.Tensor) + set_track_meta(True) def test_k(self): key = "test" rotate = RandRotate90d(keys=key, max_k=2) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate.set_random_state(234) - rotated = rotate({key: p(self.imt[0])}) + im = {key: p(self.imt[0])} + rotated = rotate(im) + test_local_inversion(rotate, rotated, im, key) expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") def test_spatial_axes(self): key = "test" rotate = RandRotate90d(keys=key, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate.set_random_state(234) - rotated = rotate({key: p(self.imt[0])}) + im = {key: p(self.imt[0])} + rotated = rotate(im) + test_local_inversion(rotate, rotated, im, key) expected = [np.rot90(channel, 0, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") def test_prob_k_spatial_axes(self): key = "test" rotate = RandRotate90d(keys=key, prob=1.0, max_k=2, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate.set_random_state(234) - rotated = rotate({key: p(self.imt[0])}) + im = {key: p(self.imt[0])} + rotated = rotate(im) expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") + test_local_inversion(rotate, rotated, im, key) def test_no_key(self): key = "unknown" diff --git a/tests/test_rand_rotated.py b/tests/test_rand_rotated.py index 464b37d925..906977f3fa 100644 --- a/tests/test_rand_rotated.py +++ b/tests/test_rand_rotated.py @@ -19,10 +19,10 @@ from monai.transforms import RandRotated from monai.utils import GridSampleMode, GridSamplePadMode -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D, test_local_inversion TEST_CASES_2D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_2D.append((p, np.pi / 2, True, "bilinear", "border", False)) TEST_CASES_2D.append((p, np.pi / 4, True, "nearest", "border", False)) TEST_CASES_2D.append((p, np.pi, False, "nearest", "zeros", True)) @@ -30,7 +30,7 @@ TEST_CASES_3D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_3D.append( (p, np.pi / 2, -np.pi / 6, (0.0, np.pi), False, "bilinear", "border", False, (1, 87, 104, 109)) ) @@ -118,8 +118,9 @@ def test_correct_results(self, im_type, degrees, keep_size, mode, padding_mode, align_corners=align_corners, dtype=np.float64, ) + im = im_type(self.imt[0]) rotate_fn.set_random_state(243) - rotated = rotate_fn({"img": im_type(self.imt[0]), "seg": im_type(self.segn[0])}) + rotated = rotate_fn({"img": im, "seg": im_type(self.segn[0])}) _order = 0 if mode == "nearest" else 1 if padding_mode == "border": @@ -132,6 +133,7 @@ def test_correct_results(self, im_type, degrees, keep_size, mode, padding_mode, expected = scipy.ndimage.rotate( self.imt[0, 0], -np.rad2deg(angle), (0, 1), not keep_size, order=_order, mode=_mode, prefilter=False ) + test_local_inversion(rotate_fn, rotated, {"img": im}, "img") for k, v in rotated.items(): rotated[k] = v.cpu() if isinstance(v, torch.Tensor) else v expected = np.stack(expected).astype(np.float32) diff --git a/tests/test_rand_scale_crop.py b/tests/test_rand_scale_crop.py index 5d6312002f..a97a77a8e6 100644 --- a/tests/test_rand_scale_crop.py +++ b/tests/test_rand_scale_crop.py @@ -15,66 +15,57 @@ from parameterized import parameterized from monai.transforms import RandScaleCrop -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose -TEST_CASE_1 = [ - {"roi_scale": [1.0, 1.0, -1.0], "random_center": True}, - np.random.randint(0, 2, size=[3, 3, 3, 4]), - (3, 3, 3, 4), +TEST_SHAPES = [ + [{"roi_scale": [1.0, 1.0, -1.0], "random_center": True}, (3, 3, 3, 4), (3, 3, 3, 4)], + [{"roi_scale": [1.0, 1.0, 1.0], "random_center": False}, (3, 3, 3, 3), (3, 3, 3, 3)], ] -TEST_CASE_2 = [ - {"roi_scale": [1.0, 1.0, 1.0], "random_center": False}, - np.random.randint(0, 2, size=[3, 3, 3, 3]), - (3, 3, 3, 3), +TEST_VALUES = [ + [ + {"roi_scale": [0.6, 0.6], "random_center": False}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + ] ] -TEST_CASE_3 = [ - {"roi_scale": [0.6, 0.6], "random_center": False}, - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), +TEST_RANDOM_SHAPES = [ + [ + {"roi_scale": [0.75, 0.6, 0.5], "max_roi_scale": [1.0, -1.0, 0.6], "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 3, 4, 3), + ], + [{"roi_scale": 0.6, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, (1, 4, 5, 6), (1, 3, 4, 4)], + [{"roi_scale": 0.2, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, (1, 4, 5, 6), (1, 3, 2, 4)], ] -TEST_CASE_4 = [ - {"roi_scale": [0.75, 0.6, 0.5], "max_roi_scale": [1.0, -1.0, 0.6], "random_center": True, "random_size": True}, - np.random.randint(0, 2, size=[1, 4, 5, 6]), - (1, 3, 4, 3), -] - -TEST_CASE_5 = [ - {"roi_scale": 0.6, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, - np.random.randint(0, 2, size=[1, 4, 5, 6]), - (1, 3, 4, 4), -] - -TEST_CASE_6 = [ - {"roi_scale": 0.2, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, - np.random.randint(0, 2, size=[1, 4, 5, 6]), - (1, 3, 2, 4), -] +class TestRandScaleCrop(CropTest): + Cropper = RandScaleCrop -class TestRandScaleCrop(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) - def test_shape(self, input_param, input_data, expected_shape): - for p in TEST_NDARRAYS: - result = RandScaleCrop(**input_param)(p(input_data)) - self.assertTupleEqual(result.shape, expected_shape) + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_3]) + @parameterized.expand(TEST_VALUES) def test_value(self, input_param, input_data): - for p in TEST_NDARRAYS: - cropper = RandScaleCrop(**input_param) - result = cropper(p(input_data)) - roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] - assert_allclose(result, input_data[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test=False) + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = RandScaleCrop(**input_param) + result = cropper(im_type(input_data)) + roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] + assert_allclose(result, input_data[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test="tensor") - @parameterized.expand([TEST_CASE_4, TEST_CASE_5, TEST_CASE_6]) - def test_random_shape(self, input_param, input_data, expected_shape): - for p in TEST_NDARRAYS: - cropper = RandScaleCrop(**input_param) - cropper.set_random_state(seed=123) - result = cropper(p(input_data)) - self.assertTupleEqual(result.shape, expected_shape) + @parameterized.expand(TEST_RANDOM_SHAPES) + def test_random_shape(self, input_param, input_shape, expected_shape): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = RandScaleCrop(**input_param) + cropper.set_random_state(seed=123) + input_data = im_type(np.random.randint(0, 2, input_shape)) + result = cropper(input_data) + self.assertTupleEqual(result.shape, expected_shape) if __name__ == "__main__": diff --git a/tests/test_rand_scale_cropd.py b/tests/test_rand_scale_cropd.py index 5e833fef98..dd92783766 100644 --- a/tests/test_rand_scale_cropd.py +++ b/tests/test_rand_scale_cropd.py @@ -15,74 +15,77 @@ from parameterized import parameterized from monai.transforms import RandScaleCropd -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose -TEST_CASE_1 = [ - {"keys": "img", "roi_scale": [1.0, 1.0, -1.0], "random_center": True}, - {"img": np.random.randint(0, 2, size=[3, 3, 3, 4])}, - (3, 3, 3, 4), +TEST_SHAPES = [ + [{"keys": "img", "roi_scale": [1.0, 1.0, -1.0], "random_center": True}, (3, 3, 3, 4), (3, 3, 3, 4)], + [ + # test `allow_missing_keys` with key "label" + {"keys": ["label", "img"], "roi_scale": [1.0, 1.0, 1.0], "random_center": False, "allow_missing_keys": True}, + (3, 3, 3, 3), + (3, 3, 3, 3), + ], ] -TEST_CASE_2 = [ - # test `allow_missing_keys` with key "label" - {"keys": ["label", "img"], "roi_scale": [1.0, 1.0, 1.0], "random_center": False, "allow_missing_keys": True}, - {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, - (3, 3, 3, 3), +TEST_VALUES = [ + [ + {"keys": "img", "roi_scale": [0.6, 0.6], "random_center": False}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + ] ] -TEST_CASE_3 = [ - {"keys": "img", "roi_scale": [0.6, 0.6], "random_center": False}, - {"img": np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]])}, +TEST_RANDOM_SHAPES = [ + [ + { + "keys": "img", + "roi_scale": [0.75, 0.6, 0.5], + "max_roi_scale": [1.0, -1.0, 0.6], + "random_center": True, + "random_size": True, + }, + (1, 4, 5, 6), + (1, 3, 4, 3), + ], + [ + {"keys": "img", "roi_scale": 0.6, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 3, 4, 4), + ], + [ + {"keys": "img", "roi_scale": 0.2, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 3, 2, 4), + ], ] -TEST_CASE_4 = [ - { - "keys": "img", - "roi_scale": [0.75, 0.6, 0.5], - "max_roi_scale": [1.0, -1.0, 0.6], - "random_center": True, - "random_size": True, - }, - {"img": np.random.randint(0, 2, size=[1, 4, 5, 6])}, - (1, 3, 4, 3), -] - -TEST_CASE_5 = [ - {"keys": "img", "roi_scale": 0.6, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, - {"img": np.random.randint(0, 2, size=[1, 4, 5, 6])}, - (1, 3, 4, 4), -] - -TEST_CASE_6 = [ - {"keys": "img", "roi_scale": 0.2, "max_roi_scale": 0.8, "random_center": True, "random_size": True}, - {"img": np.random.randint(0, 2, size=[1, 4, 5, 6])}, - (1, 3, 2, 4), -] +class TestRandScaleCropd(CropTest): + Cropper = RandScaleCropd -class TestRandScaleCropd(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) - def test_shape(self, input_param, input_data, expected_shape): - result = RandScaleCropd(**input_param)(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_3]) - def test_value(self, input_param, input_data): - for p in TEST_NDARRAYS: - cropper = RandScaleCropd(**input_param) - input_data["img"] = p(input_data["img"]) - result = cropper(input_data) - roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] - assert_allclose( - result["img"], input_data["img"][:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test=False - ) + @parameterized.expand(TEST_VALUES) + def test_value(self, input_param, input_im): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = self.Cropper(**input_param) + input_data = {"img": im_type(input_im)} + result = cropper(input_data)["img"] + roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper.cropper._size] + assert_allclose(result, input_im[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test="tensor") - @parameterized.expand([TEST_CASE_4, TEST_CASE_5, TEST_CASE_6]) - def test_random_shape(self, input_param, input_data, expected_shape): - cropper = RandScaleCropd(**input_param) - cropper.set_random_state(seed=123) - result = cropper(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + @parameterized.expand(TEST_RANDOM_SHAPES) + def test_random_shape(self, input_param, input_shape, expected_shape): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = self.Cropper(**input_param) + cropper.set_random_state(seed=123) + input_data = {"img": im_type(np.random.randint(0, 2, input_shape))} + result = cropper(input_data)["img"] + self.assertTupleEqual(result.shape, expected_shape) if __name__ == "__main__": diff --git a/tests/test_rand_scale_intensity.py b/tests/test_rand_scale_intensity.py index 5aa5c7b964..b0999a82a5 100644 --- a/tests/test_rand_scale_intensity.py +++ b/tests/test_rand_scale_intensity.py @@ -12,22 +12,24 @@ import unittest import numpy as np +from parameterized import parameterized from monai.transforms import RandScaleIntensity from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose class TestRandScaleIntensity(NumpyImageTestCase2D): - def test_value(self): - for p in TEST_NDARRAYS: - scaler = RandScaleIntensity(factors=0.5, prob=1.0) - scaler.set_random_state(seed=0) - result = scaler(p(self.imt)) - np.random.seed(0) - # simulate the randomize() of transform - np.random.random() - expected = p((self.imt * (1 + np.random.uniform(low=-0.5, high=0.5))).astype(np.float32)) - assert_allclose(result, p(expected), rtol=1e-7, atol=0) + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_value(self, p): + scaler = RandScaleIntensity(factors=0.5, prob=1.0) + scaler.set_random_state(seed=0) + im = p(self.imt) + result = scaler(im) + np.random.seed(0) + # simulate the randomize() of transform + np.random.random() + expected = p((self.imt * (1 + np.random.uniform(low=-0.5, high=0.5))).astype(np.float32)) + assert_allclose(result, p(expected), rtol=1e-7, atol=0, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_scale_intensityd.py b/tests/test_rand_scale_intensityd.py index 655bd88ee0..d548ee34d6 100644 --- a/tests/test_rand_scale_intensityd.py +++ b/tests/test_rand_scale_intensityd.py @@ -28,7 +28,7 @@ def test_value(self): # simulate the randomize function of transform np.random.random() expected = (self.imt * (1 + np.random.uniform(low=-0.5, high=0.5))).astype(np.float32) - assert_allclose(result[key], p(expected)) + assert_allclose(result[key], p(expected), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_shift_intensity.py b/tests/test_rand_shift_intensity.py index b4f32a385a..d5ad083d33 100644 --- a/tests/test_rand_shift_intensity.py +++ b/tests/test_rand_shift_intensity.py @@ -12,21 +12,24 @@ import unittest import numpy as np +from parameterized import parameterized from monai.transforms import RandShiftIntensity -from tests.utils import NumpyImageTestCase2D +from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose class TestRandShiftIntensity(NumpyImageTestCase2D): - def test_value(self): + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_value(self, p): shifter = RandShiftIntensity(offsets=1.0, prob=1.0) shifter.set_random_state(seed=0) - result = shifter(self.imt, factor=1.0) + im = p(self.imt) + result = shifter(im, factor=1.0) np.random.seed(0) # simulate the randomize() of transform np.random.random() expected = self.imt + np.random.uniform(low=-1.0, high=1.0) - np.testing.assert_allclose(result, expected) + assert_allclose(result, expected, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_shift_intensityd.py b/tests/test_rand_shift_intensityd.py index 4d05149e3c..1a8356c2c9 100644 --- a/tests/test_rand_shift_intensityd.py +++ b/tests/test_rand_shift_intensityd.py @@ -29,7 +29,7 @@ def test_value(self): # simulate the randomize() of transform np.random.random() expected = self.imt + np.random.uniform(low=-1.0, high=1.0) - assert_allclose(result[key], p(expected)) + assert_allclose(result[key], p(expected), type_test="tensor") def test_factor(self): key = "img" diff --git a/tests/test_rand_spatial_crop.py b/tests/test_rand_spatial_crop.py index 8f4bb0fffa..383ea8a1cb 100644 --- a/tests/test_rand_spatial_crop.py +++ b/tests/test_rand_spatial_crop.py @@ -15,60 +15,57 @@ from parameterized import parameterized from monai.transforms import RandSpatialCrop -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose -TEST_CASE_0 = [ - {"roi_size": [3, 3, -1], "random_center": True}, - np.random.randint(0, 2, size=[3, 3, 3, 4]), - (3, 3, 3, 4), +TEST_SHAPES = [ + [{"roi_size": [3, 3, -1], "random_center": True}, (3, 3, 3, 4), (3, 3, 3, 4)], + [{"roi_size": [3, 3, 3], "random_center": True}, (3, 3, 3, 3), (3, 3, 3, 3)], + [{"roi_size": [3, 3, 3], "random_center": False}, (3, 3, 3, 3), (3, 3, 3, 3)], ] -TEST_CASE_1 = [{"roi_size": [3, 3, 3], "random_center": True}, np.random.randint(0, 2, size=[3, 3, 3, 3]), (3, 3, 3, 3)] - -TEST_CASE_2 = [ - {"roi_size": [3, 3, 3], "random_center": False}, - np.random.randint(0, 2, size=[3, 3, 3, 3]), - (3, 3, 3, 3), -] - -TEST_CASE_3 = [ - {"roi_size": [3, 3], "random_center": False}, - np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), +TEST_VALUES = [ + [ + {"roi_size": [3, 3], "random_center": False}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + ] ] -TEST_CASE_4 = [ - {"roi_size": [3, 3, 3], "max_roi_size": [5, -1, 4], "random_center": True, "random_size": True}, - np.random.randint(0, 2, size=[1, 4, 5, 6]), - (1, 4, 4, 3), +TEST_RANDOM_SHAPES = [ + [ + {"roi_size": [3, 3, 3], "max_roi_size": [5, -1, 4], "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 4, 4, 3), + ], + [{"roi_size": 3, "max_roi_size": 4, "random_center": True, "random_size": True}, (1, 4, 5, 6), (1, 3, 4, 3)], ] -TEST_CASE_5 = [ - {"roi_size": 3, "max_roi_size": 4, "random_center": True, "random_size": True}, - np.random.randint(0, 2, size=[1, 4, 5, 6]), - (1, 3, 4, 3), -] +class TestRandSpatialCrop(CropTest): + Cropper = RandSpatialCrop -class TestRandSpatialCrop(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2]) - def test_shape(self, input_param, input_data, expected_shape): - result = RandSpatialCrop(**input_param)(input_data) - self.assertTupleEqual(result.shape, expected_shape) + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_3]) + @parameterized.expand(TEST_VALUES) def test_value(self, input_param, input_data): - for p in TEST_NDARRAYS: - cropper = RandSpatialCrop(**input_param) - result = cropper(p(input_data)) - roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] - assert_allclose(result, input_data[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test=False) + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = RandSpatialCrop(**input_param) + result = cropper(im_type(input_data)) + roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] + assert_allclose(result, input_data[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test="tensor") - @parameterized.expand([TEST_CASE_4, TEST_CASE_5]) - def test_random_shape(self, input_param, input_data, expected_shape): - cropper = RandSpatialCrop(**input_param) - cropper.set_random_state(seed=123) - result = cropper(input_data) - self.assertTupleEqual(result.shape, expected_shape) + @parameterized.expand(TEST_RANDOM_SHAPES) + def test_random_shape(self, input_param, input_shape, expected_shape): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = RandSpatialCrop(**input_param) + cropper.set_random_state(seed=123) + input_data = im_type(np.random.randint(0, 2, input_shape)) + result = cropper(input_data) + self.assertTupleEqual(result.shape, expected_shape) if __name__ == "__main__": diff --git a/tests/test_rand_spatial_crop_samples.py b/tests/test_rand_spatial_crop_samples.py index 18fdf38773..fd905a6dae 100644 --- a/tests/test_rand_spatial_crop_samples.py +++ b/tests/test_rand_spatial_crop_samples.py @@ -15,11 +15,12 @@ from parameterized import parameterized from monai.transforms import RandSpatialCropSamples -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TEST_CASE_1 = [ {"roi_size": [3, 3, 3], "num_samples": 4, "random_center": True, "random_size": False}, - np.arange(192).reshape(3, 4, 4, 4), + (3, 4, 4, 4), [(3, 3, 3, 3), (3, 3, 3, 3), (3, 3, 3, 3), (3, 3, 3, 3)], np.array( [ @@ -44,7 +45,7 @@ TEST_CASE_2 = [ {"roi_size": [3, 3, 3], "num_samples": 8, "random_center": False, "random_size": True}, - np.arange(192).reshape(3, 4, 4, 4), + (3, 4, 4, 4), [(3, 4, 4, 3), (3, 4, 3, 3), (3, 3, 4, 4), (3, 4, 4, 4), (3, 3, 3, 4), (3, 3, 3, 3), (3, 3, 3, 3), (3, 3, 3, 3)], np.array( [ @@ -67,19 +68,32 @@ ), ] +TEST_INVERSE_LIST = [ + [(1, 2, 2), {"roi_size": (1, 1), "num_samples": 4, "random_size": False}], + [(1, 3, 2), {"roi_size": (1, 1), "num_samples": 100, "random_size": False}], + [(3, 10, 11, 12), {"roi_size": (3, 5, 4), "num_samples": 7, "random_size": False}], + [(3, 10, 11, 12), {"roi_size": (10, 11, 12), "num_samples": 3, "random_size": False}], + [(3, 10, 11, 12), {"roi_size": (3, 4, 5), "num_samples": 100, "random_size": False}], +] + + +class TestRandSpatialCropSamples(CropTest): + Cropper = RandSpatialCropSamples -class TestRandSpatialCropSamples(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) - def test_shape(self, input_param, input_data, expected_shape, expected_last_item): - for p in TEST_NDARRAYS: + def test_shape(self, input_param, input_shape, expected_shape, expected_last_item): + input_data = np.arange(192).reshape(*input_shape) + + for p in TEST_NDARRAYS_ALL: xform = RandSpatialCropSamples(**input_param) xform.set_random_state(1234) result = xform(p(input_data)) np.testing.assert_equal(len(result), input_param["num_samples"]) - for item, expected in zip(result, expected_shape): + for i, (item, expected) in enumerate(zip(result, expected_shape)): self.assertTupleEqual(item.shape, expected) - assert_allclose(result[-1], expected_last_item, type_test=False) + self.assertEqual(item.meta["patch_index"], i) + assert_allclose(result[-1], expected_last_item, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_spatial_crop_samplesd.py b/tests/test_rand_spatial_crop_samplesd.py index 0891068488..4da438d2a0 100644 --- a/tests/test_rand_spatial_crop_samplesd.py +++ b/tests/test_rand_spatial_crop_samplesd.py @@ -14,46 +14,45 @@ import numpy as np from parameterized import parameterized -from monai.transforms import Compose, RandSpatialCropSamplesd, ToTensord -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS, assert_allclose +from monai.transforms import Compose, DivisiblePadd, RandSpatialCropSamplesd +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TEST_CASE_1 = [ {"keys": ["img", "seg"], "num_samples": 4, "roi_size": [2, 2, 2], "random_center": True}, {"img": np.arange(81).reshape(3, 3, 3, 3), "seg": np.arange(81, 0, -1).reshape(3, 3, 3, 3)}, - [(3, 3, 3, 2), (3, 2, 2, 2), (3, 3, 3, 2), (3, 3, 2, 2)], + [(3, 2, 2, 2), (3, 2, 3, 3), (3, 2, 3, 2), (3, 2, 3, 2)], { "img": np.array( [ - [[[0, 1], [3, 4]], [[9, 10], [12, 13]], [[18, 19], [21, 22]]], - [[[27, 28], [30, 31]], [[36, 37], [39, 40]], [[45, 46], [48, 49]]], - [[[54, 55], [57, 58]], [[63, 64], [66, 67]], [[72, 73], [75, 76]]], + [[[1, 2], [4, 5], [7, 8]], [[10, 11], [13, 14], [16, 17]]], + [[[28, 29], [31, 32], [34, 35]], [[37, 38], [40, 41], [43, 44]]], + [[[55, 56], [58, 59], [61, 62]], [[64, 65], [67, 68], [70, 71]]], ] ), "seg": np.array( [ - [[[81, 80], [78, 77]], [[72, 71], [69, 68]], [[63, 62], [60, 59]]], - [[[54, 53], [51, 50]], [[45, 44], [42, 41]], [[36, 35], [33, 32]]], - [[[27, 26], [24, 23]], [[18, 17], [15, 14]], [[9, 8], [6, 5]]], + [[[80, 79], [77, 76], [74, 73]], [[71, 70], [68, 67], [65, 64]]], + [[[53, 52], [50, 49], [47, 46]], [[44, 43], [41, 40], [38, 37]]], + [[[26, 25], [23, 22], [20, 19]], [[17, 16], [14, 13], [11, 10]]], ] ), }, ] TEST_CASE_2 = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASE_2.append( [ {"keys": ["img", "seg"], "num_samples": 8, "roi_size": [2, 2, 3], "random_center": False}, {"img": p(np.arange(81).reshape(3, 3, 3, 3)), "seg": p(np.arange(81, 0, -1).reshape(3, 3, 3, 3))}, [ - (3, 3, 3, 3), - (3, 2, 3, 3), (3, 2, 2, 3), - (3, 2, 3, 3), + (3, 2, 2, 3), (3, 3, 3, 3), + (3, 2, 3, 3), (3, 3, 3, 3), - (3, 2, 2, 3), + (3, 2, 3, 3), + (3, 2, 3, 3), (3, 3, 2, 3), ], { @@ -90,10 +89,10 @@ def test_shape(self, input_param, input_data, expected_shape, expected_last): self.assertTupleEqual(item["img"].shape, expected) self.assertTupleEqual(item["seg"].shape, expected) for i, item in enumerate(result): - self.assertEqual(item[PostFix.meta("img")]["patch_index"], i) - self.assertEqual(item[PostFix.meta("seg")]["patch_index"], i) - assert_allclose(item["img"], expected_last["img"], type_test=True) - assert_allclose(item["seg"], expected_last["seg"], type_test=True) + self.assertEqual(item["img"].meta["patch_index"], i) + self.assertEqual(item["seg"].meta["patch_index"], i) + assert_allclose(item["img"], expected_last["img"], type_test=False) + assert_allclose(item["seg"], expected_last["seg"], type_test=False) def test_deep_copy(self): data = {"img": np.ones((1, 10, 11, 12))} @@ -101,11 +100,11 @@ def test_deep_copy(self): sampler = RandSpatialCropSamplesd( keys=["img"], roi_size=(3, 3, 3), num_samples=num_samples, random_center=True, random_size=False ) - transform = Compose([ToTensord(keys="img"), sampler]) + transform = Compose([DivisiblePadd(keys="img", k=5), sampler]) samples = transform(data) self.assertEqual(len(samples), num_samples) for sample in samples: - self.assertEqual(len(sample["img_transforms"]), len(transform)) + self.assertEqual(len(sample["img"].applied_operations), len(transform)) if __name__ == "__main__": diff --git a/tests/test_rand_spatial_cropd.py b/tests/test_rand_spatial_cropd.py index 9e6e86eea2..1b256959c6 100644 --- a/tests/test_rand_spatial_cropd.py +++ b/tests/test_rand_spatial_cropd.py @@ -15,65 +15,62 @@ from parameterized import parameterized from monai.transforms import RandSpatialCropd -from tests.utils import TEST_NDARRAYS +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose -TEST_CASE_0 = [ - {"keys": "img", "roi_size": [3, 3, -1], "random_center": True}, - {"img": np.random.randint(0, 2, size=[3, 3, 3, 5])}, - (3, 3, 3, 5), +TEST_SHAPES = [ + [{"keys": "img", "roi_size": [3, 3, -1], "random_center": True}, (3, 3, 3, 5), (3, 3, 3, 5)], + [{"keys": "img", "roi_size": [3, 3, 3], "random_center": True}, (3, 3, 3, 3), (3, 3, 3, 3)], + [{"keys": "img", "roi_size": [3, 3, 3], "random_center": False}, (3, 3, 3, 3), (3, 3, 3, 3)], ] -TEST_CASE_1 = [ - {"keys": "img", "roi_size": [3, 3, 3], "random_center": True}, - {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, - (3, 3, 3, 3), +TEST_VALUES = [ + [ + {"keys": "img", "roi_size": [3, 3], "random_center": False}, + np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]]), + ] ] -TEST_CASE_2 = [ - {"keys": "img", "roi_size": [3, 3, 3], "random_center": False}, - {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])}, - (3, 3, 3, 3), +TEST_RANDOM_SHAPES = [ + [ + {"keys": "img", "roi_size": [3, 3, 3], "max_roi_size": [5, -1, 4], "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 4, 4, 3), + ], + [ + {"keys": "img", "roi_size": 3, "max_roi_size": 4, "random_center": True, "random_size": True}, + (1, 4, 5, 6), + (1, 3, 4, 3), + ], ] -TEST_CASE_3 = [ - {"keys": "img", "roi_size": [3, 3], "random_center": False}, - {"img": np.array([[[0, 0, 0, 0, 0], [0, 1, 2, 1, 0], [0, 2, 3, 2, 0], [0, 1, 2, 1, 0], [0, 0, 0, 0, 0]]])}, -] - -TEST_CASE_4 = [ - {"keys": "img", "roi_size": [3, 3, 3], "max_roi_size": [5, -1, 4], "random_center": True, "random_size": True}, - {"img": np.random.randint(0, 2, size=[1, 4, 5, 6])}, - (1, 4, 4, 3), -] - -TEST_CASE_5 = [ - {"keys": "img", "roi_size": 3, "max_roi_size": 4, "random_center": True, "random_size": True}, - {"img": np.random.randint(0, 2, size=[1, 4, 5, 6])}, - (1, 3, 4, 3), -] +class TestRandSpatialCropd(CropTest): + Cropper = RandSpatialCropd -class TestRandSpatialCropd(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2]) - def test_shape(self, input_param, input_data, expected_shape): - result = RandSpatialCropd(**input_param)(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + @parameterized.expand(TEST_SHAPES) + def test_shape(self, input_param, input_shape, expected_shape): + self.crop_test(input_param, input_shape, expected_shape) - @parameterized.expand([TEST_CASE_3]) - def test_value(self, input_param, input_data): - cropper = RandSpatialCropd(**input_param) - result = cropper(input_data) - roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper._size] - np.testing.assert_allclose(result["img"], input_data["img"][:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]]) + @parameterized.expand(TEST_VALUES) + def test_value(self, input_param, input_im): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = self.Cropper(**input_param) + input_data = {"img": im_type(input_im)} + result = cropper(input_data)["img"] + roi = [(2 - i // 2, 2 + i - i // 2) for i in cropper.cropper._size] + assert_allclose(result, input_im[:, roi[0][0] : roi[0][1], roi[1][0] : roi[1][1]], type_test="tensor") - @parameterized.expand([TEST_CASE_4, TEST_CASE_5]) - def test_random_shape(self, input_param, input_data, expected_shape): - for p in TEST_NDARRAYS: - cropper = RandSpatialCropd(**input_param) - cropper.set_random_state(seed=123) - input_data["img"] = p(input_data["img"]) - result = cropper(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + @parameterized.expand(TEST_RANDOM_SHAPES) + def test_random_shape(self, input_param, input_shape, expected_shape): + for im_type in TEST_NDARRAYS_ALL: + with self.subTest(im_type=im_type): + cropper = self.Cropper(**input_param) + cropper.set_random_state(seed=123) + input_data = {"img": im_type(np.random.randint(0, 2, input_shape))} + result = cropper(input_data)["img"] + self.assertTupleEqual(result.shape, expected_shape) if __name__ == "__main__": diff --git a/tests/test_rand_std_shift_intensity.py b/tests/test_rand_std_shift_intensity.py index fdf386fee4..b26f5ef096 100644 --- a/tests/test_rand_std_shift_intensity.py +++ b/tests/test_rand_std_shift_intensity.py @@ -12,25 +12,25 @@ import unittest import numpy as np -import torch +from parameterized import parameterized from monai.transforms import RandStdShiftIntensity -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D +from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose class TestRandStdShiftIntensity(NumpyImageTestCase2D): - def test_value(self): - for p in TEST_NDARRAYS: - np.random.seed(0) - # simulate the randomize() of transform - np.random.random() - factor = np.random.uniform(low=-1.0, high=1.0) - offset = factor * np.std(self.imt) - expected = p(self.imt + offset) - shifter = RandStdShiftIntensity(factors=1.0, prob=1.0) - shifter.set_random_state(seed=0) - result = shifter(p(self.imt)) - torch.testing.assert_allclose(result, expected, atol=0, rtol=1e-5) + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_value(self, p): + np.random.seed(0) + # simulate the randomize() of transform + np.random.random() + factor = np.random.uniform(low=-1.0, high=1.0) + offset = factor * np.std(self.imt) + expected = p(self.imt + offset) + shifter = RandStdShiftIntensity(factors=1.0, prob=1.0) + shifter.set_random_state(seed=0) + result = shifter(p(self.imt)) + assert_allclose(result, expected, atol=0, rtol=1e-5, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_std_shift_intensityd.py b/tests/test_rand_std_shift_intensityd.py index e98d1e3ad3..bbbed053ad 100644 --- a/tests/test_rand_std_shift_intensityd.py +++ b/tests/test_rand_std_shift_intensityd.py @@ -12,10 +12,9 @@ import unittest import numpy as np -import torch from monai.transforms import RandStdShiftIntensityd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D +from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose class TestRandStdShiftIntensityd(NumpyImageTestCase2D): @@ -30,9 +29,7 @@ def test_value(self): shifter = RandStdShiftIntensityd(keys=[key], factors=1.0, prob=1.0) shifter.set_random_state(seed=0) result = shifter({key: p(self.imt)})[key] - if isinstance(result, torch.Tensor): - result = result.cpu() - np.testing.assert_allclose(result, expected, rtol=1e-5) + assert_allclose(result, expected, rtol=1e-5, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rand_weighted_crop.py b/tests/test_rand_weighted_crop.py index dae7f05016..53913ce987 100644 --- a/tests/test_rand_weighted_crop.py +++ b/tests/test_rand_weighted_crop.py @@ -12,11 +12,11 @@ import unittest import numpy as np -import torch from parameterized.parameterized import parameterized from monai.transforms.croppad.array import RandWeightedCrop -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D, assert_allclose +from tests.croppers import CropTest +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D, assert_allclose def get_data(ndim): @@ -30,8 +30,8 @@ def get_data(ndim): TESTS = [] -for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: + for q in TEST_NDARRAYS_ALL: im = SEG1_2D weight = np.zeros_like(im) weight[0, 30, 17] = 1.1 @@ -148,7 +148,9 @@ def get_data(ndim): ) -class TestRandWeightedCrop(unittest.TestCase): +class TestRandWeightedCrop(CropTest): + Cropper = RandWeightedCrop + @parameterized.expand(TESTS) def test_rand_weighted_crop(self, _, input_params, img, weight, expected_shape, expected_vals): crop = RandWeightedCrop(**input_params) @@ -161,10 +163,8 @@ def test_rand_weighted_crop(self, _, input_params, img, weight, expected_shape, # if desired ROI is larger than image, check image is unchanged if all(s >= i for i, s in zip(img.shape[1:], input_params["spatial_size"])): for res in result: - self.assertEqual(type(img), type(res)) - if isinstance(img, torch.Tensor): - self.assertEqual(res.device, img.device) - assert_allclose(res, img) + assert_allclose(res, img, type_test="tensor") + self.assertEqual(len(res.applied_operations), 1) if __name__ == "__main__": diff --git a/tests/test_rand_weighted_cropd.py b/tests/test_rand_weighted_cropd.py index a357398f1c..b3fc92b445 100644 --- a/tests/test_rand_weighted_cropd.py +++ b/tests/test_rand_weighted_cropd.py @@ -12,180 +12,144 @@ import unittest import numpy as np +from parameterized import parameterized from monai.transforms.croppad.dictionary import RandWeightedCropd -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D, assert_allclose - - -class TestRandWeightedCrop(NumpyImageTestCase2D): - def test_rand_weighted_crop_small_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.seg1[0] - n_samples = 3 - crop = RandWeightedCropd("img", "w", (10, 12), n_samples) - weight = np.zeros_like(img) - weight[0, 30, 17] = 1.1 - weight[0, 40, 31] = 1 - weight[0, 80, 21] = 1 - crop.set_random_state(10) - d = {"img": p(img), "w": q(weight)} - result = crop(d) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 10, 12)) - for c, e in zip(crop.centers, [[80, 21], [30, 17], [40, 31]]): - assert_allclose(c, e, type_test=False) - - def test_rand_weighted_crop_default_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.imt[0] - n_samples = 3 - crop = RandWeightedCropd("im", "weight", (10, -1), n_samples, "coords") - weight = np.zeros_like(img) - weight[0, 30, 17] = 1.1 - weight[0, 40, 31] = 1 - weight[0, 80, 21] = 1 - crop.set_random_state(10) - data = {"im": p(img), "weight": q(weight), "others": np.nan} - result = crop(data) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["im"].shape, (1, 10, 64)) - for c, e in zip(crop.centers, [[14, 32], [105, 32], [20, 32]]): - assert_allclose(c, e, type_test=False) - assert_allclose(result[1]["coords"], [105, 32], type_test=False) - - def test_rand_weighted_crop_large_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.segn[0] - n_samples = 3 - crop = RandWeightedCropd(("img", "seg"), "weight", (10000, 400), n_samples, "location") - weight = np.zeros_like(img) - weight[0, 30, 17] = 1.1 - weight[0, 10, 1] = 1 - crop.set_random_state(10) - data = {"img": p(img), "seg": p(self.imt[0]), "weight": q(weight)} - result = crop(data) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 128, 64)) - np.testing.assert_allclose(result[0]["seg"].shape, (1, 128, 64)) - for c, e in zip(crop.centers, [[64, 32], [64, 32], [64, 32]]): - assert_allclose(c, e, type_test=False) - assert_allclose(result[1]["location"], [64, 32], type_test=False) - - def test_rand_weighted_crop_bad_w(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.imt[0] - n_samples = 3 - crop = RandWeightedCropd(("img", "seg"), "w", (20, 40), n_samples) - weight = np.zeros_like(img) - weight[0, 30, 17] = np.inf - weight[0, 10, 1] = -np.inf - weight[0, 10, 20] = -np.nan - crop.set_random_state(10) - result = crop({"img": p(img), "seg": p(self.segn[0]), "w": q(weight)}) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 20, 40)) - np.testing.assert_allclose(result[0]["seg"].shape, (1, 20, 40)) - for c, e in zip(crop.centers, [[63, 37], [31, 43], [66, 20]]): - assert_allclose(c, e, type_test=False) - - -class TestRandWeightedCrop3D(NumpyImageTestCase3D): - def test_rand_weighted_crop_small_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.seg1[0] - n_samples = 3 - crop = RandWeightedCropd("img", "w", (8, 10, 12), n_samples) - weight = np.zeros_like(img) - weight[0, 5, 30, 17] = 1.1 - weight[0, 8, 40, 31] = 1 - weight[0, 11, 23, 21] = 1 - crop.set_random_state(10) - result = crop({"img": p(img), "w": q(weight)}) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 8, 10, 12)) - for c, e in zip(crop.centers, [[11, 23, 21], [5, 30, 17], [8, 40, 31]]): - assert_allclose(c, e, type_test=False) - - def test_rand_weighted_crop_default_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.imt[0] - n_samples = 3 - crop = RandWeightedCropd(("img", "seg"), "w", (10, -1, -1), n_samples) - weight = np.zeros_like(img) - weight[0, 7, 17] = 1.1 - weight[0, 13, 31] = 1.1 - weight[0, 24, 21] = 1 - crop.set_random_state(10) - result = crop({"img": p(img), "seg": p(self.segn[0]), "w": q(weight)}) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 10, 64, 80)) - np.testing.assert_allclose(result[0]["seg"].shape, (1, 10, 64, 80)) - for c, e in zip(crop.centers, [[14, 32, 40], [41, 32, 40], [20, 32, 40]]): - assert_allclose(c, e, type_test=False) - - def test_rand_weighted_crop_large_roi(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.segn[0] - n_samples = 3 - crop = RandWeightedCropd("img", "w", (10000, 400, 80), n_samples) - weight = np.zeros_like(img) - weight[0, 30, 17, 20] = 1.1 - weight[0, 10, 1, 17] = 1 - crop.set_random_state(10) - result = crop({"img": p(img), "w": q(weight)}) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 48, 64, 80)) - for c, e in zip(crop.centers, [[24, 32, 40], [24, 32, 40], [24, 32, 40]]): - assert_allclose(c, e, type_test=False) - - def test_rand_weighted_crop_bad_w(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.imt[0] - n_samples = 3 - crop = RandWeightedCropd(("img", "seg"), "w", (48, 64, 80), n_samples) - weight = np.zeros_like(img) - weight[0, 30, 17] = np.inf - weight[0, 10, 1] = -np.inf - weight[0, 10, 20] = -np.nan - crop.set_random_state(10) - result = crop({"img": p(img), "seg": p(self.segn[0]), "w": q(weight)}) - self.assertTrue(len(result) == n_samples) - np.testing.assert_allclose(result[0]["img"].shape, (1, 48, 64, 80)) - np.testing.assert_allclose(result[0]["seg"].shape, (1, 48, 64, 80)) - for c, e in zip(crop.centers, [[24, 32, 40], [24, 32, 40], [24, 32, 40]]): - assert_allclose(c, e, type_test=False) - - def test_rand_weighted_crop_patch_index(self): - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: - img = self.imt[0] - n_samples = 3 - crop = RandWeightedCropd(("img", "seg"), "w", (10, -1, -1), n_samples) - weight = np.zeros_like(img) - weight[0, 7, 17] = 1.1 - weight[0, 13, 31] = 1.1 - weight[0, 24, 21] = 1 - crop.set_random_state(10) - result = crop( - {"img": p(img), "seg": p(self.segn[0]), "w": q(weight), PostFix.meta("img"): {"affine": None}} - ) - self.assertTrue(len(result) == n_samples) - for c, e in zip(crop.centers, [[14, 32, 40], [41, 32, 40], [20, 32, 40]]): - assert_allclose(c, e, type_test=False) - for i in range(n_samples): - np.testing.assert_allclose(result[i]["img"].shape, (1, 10, 64, 80)) - np.testing.assert_allclose(result[i]["seg"].shape, (1, 10, 64, 80)) - np.testing.assert_allclose(result[i][PostFix.meta("img")]["patch_index"], i) - np.testing.assert_allclose(result[i][PostFix.meta("seg")]["patch_index"], i) +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D + + +def get_data(ndim): + im_gen = NumpyImageTestCase2D() if ndim == 2 else NumpyImageTestCase3D() + im_gen.setUp() + return im_gen.imt[0], im_gen.seg1[0], im_gen.segn[0] + + +IMT_2D, SEG1_2D, SEGN_2D = get_data(ndim=2) +IMT_3D, SEG1_3D, SEGN_3D = get_data(ndim=3) + +TESTS = [] +for p in TEST_NDARRAYS_ALL: + for q in TEST_NDARRAYS_ALL: + im = IMT_2D + weight = np.zeros_like(im) + weight[0, 30, 17] = 1.1 + weight[0, 40, 31] = 1 + weight[0, 80, 21] = 1 + TESTS.append( + [ + "small roi 2d", + dict(keys="img", w_key="w", spatial_size=(10, 12), num_samples=3), + {"img": p(im), "w": q(weight)}, + (1, 10, 12), + [[80, 21], [30, 17], [40, 31]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 30, 17] = 1.1 + weight[0, 40, 31] = 1 + weight[0, 80, 21] = 1 + TESTS.append( + [ + "default roi 2d", + dict(keys="img", w_key="w", spatial_size=(10, -1), num_samples=3), + {"img": p(im), "w": q(weight), "others": np.nan}, + (1, 10, 64), + [[14, 32], [105, 32], [20, 32]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 30, 17] = 1.1 + weight[0, 10, 1] = 1 + TESTS.append( + [ + "large roi 2d", + dict(keys=("img", "seg"), w_key="weight", spatial_size=(10000, 400), num_samples=3), + {"img": p(im), "seg": p(SEGN_2D), "weight": q(weight)}, + (1, 128, 64), + [[64, 32], [64, 32], [64, 32]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 30, 17] = np.inf + weight[0, 10, 1] = -np.inf + weight[0, 10, 20] = -np.nan + TESTS.append( + [ + "bad w roi 2d", + dict(keys=("img", "seg"), w_key="w", spatial_size=(20, 40), num_samples=3), + {"img": p(im), "seg": p(SEGN_2D), "w": q(weight)}, + (1, 20, 40), + [[63, 37], [31, 43], [66, 20]], + ] + ) + + im = IMT_3D + weight = np.zeros_like(im) + weight[0, 5, 30, 17] = 1.1 + weight[0, 8, 40, 31] = 1 + weight[0, 11, 23, 21] = 1 + TESTS.append( + [ + "small roi 3d", + dict(keys="img", w_key="w", spatial_size=(8, 10, 12), num_samples=3), + {"img": p(im), "w": q(weight)}, + (1, 8, 10, 12), + [[11, 23, 21], [5, 30, 17], [8, 40, 31]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 5, 30, 17] = 1.1 + weight[0, 8, 40, 31] = 1 + weight[0, 11, 23, 21] = 1 + TESTS.append( + [ + "default roi 3d", + dict(keys=("img", "seg"), w_key="w", spatial_size=(10, -1, -1), num_samples=3), + {"img": p(im), "seg": p(SEGN_3D), "w": q(weight)}, + (1, 10, 64, 80), + [[14, 32, 40], [41, 32, 40], [20, 32, 40]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 30, 17, 20] = 1.1 + weight[0, 10, 1, 17] = 1 + TESTS.append( + [ + "large roi 3d", + dict(keys="img", w_key="w", spatial_size=(10000, 400, 80), num_samples=3), + {"img": p(im), "w": q(weight)}, + (1, 48, 64, 80), + [[24, 32, 40], [24, 32, 40], [24, 32, 40]], + ] + ) + + weight = np.zeros_like(im) + weight[0, 30, 17] = np.inf + weight[0, 10, 1] = -np.inf + weight[0, 10, 20] = -np.nan + TESTS.append( + [ + "bad w roi 3d", + dict(keys=("img", "seg"), w_key="w", spatial_size=(48, 64, 80), num_samples=3), + {"img": p(im), "seg": p(SEGN_3D), "w": q(weight)}, + (1, 48, 64, 80), + [[24, 32, 40], [24, 32, 40], [24, 32, 40]], + ] + ) + + +class TestRandWeightedCrop(unittest.TestCase): + @parameterized.expand(TESTS) + def test_rand_weighted_cropd(self, _, init_params, input_data, expected_shape, expected_centers): + crop = RandWeightedCropd(**init_params) + crop.set_random_state(10) + result = crop(input_data) + self.assertTrue(len(result) == init_params["num_samples"]) if __name__ == "__main__": diff --git a/tests/test_rand_zoom.py b/tests/test_rand_zoom.py index 35472024ef..fc8280490f 100644 --- a/tests/test_rand_zoom.py +++ b/tests/test_rand_zoom.py @@ -16,8 +16,8 @@ from scipy.ndimage import zoom as zoom_scipy from monai.transforms import RandZoom -from monai.utils import GridSampleMode, InterpolateMode -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from monai.utils import InterpolateMode +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion VALID_CASES = [(0.8, 1.2, "nearest", False), (0.8, 1.2, InterpolateMode.NEAREST, False)] @@ -25,23 +25,27 @@ class TestRandZoom(NumpyImageTestCase2D): @parameterized.expand(VALID_CASES) def test_correct_results(self, min_zoom, max_zoom, mode, keep_size): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: random_zoom = RandZoom(prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, mode=mode, keep_size=keep_size) random_zoom.set_random_state(1234) - zoomed = random_zoom(p(self.imt[0])) + im = p(self.imt[0]) + zoomed = random_zoom(im) + test_local_inversion(random_zoom, zoomed, im) expected = [ zoom_scipy(channel, zoom=random_zoom._zoom, mode="nearest", order=0, prefilter=False) for channel in self.imt[0] ] expected = np.stack(expected).astype(np.float32) - assert_allclose(zoomed, p(expected), atol=1.0) + assert_allclose(zoomed, p(expected), atol=1.0, type_test=False) def test_keep_size(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: im = p(self.imt[0]) random_zoom = RandZoom(prob=1.0, min_zoom=0.6, max_zoom=0.7, keep_size=True) + random_zoom.set_random_state(12) zoomed = random_zoom(im) + test_local_inversion(random_zoom, zoomed, im) self.assertTrue(np.array_equal(zoomed.shape, self.imt.shape[1:])) zoomed = random_zoom(im) self.assertTrue(np.array_equal(zoomed.shape, self.imt.shape[1:])) @@ -49,26 +53,22 @@ def test_keep_size(self): self.assertTrue(np.array_equal(zoomed.shape, self.imt.shape[1:])) @parameterized.expand( - [ - ("no_min_zoom", None, 1.1, "bilinear", TypeError), - ("invalid_mode", 0.9, 1.1, "s", ValueError), - ("invalid_mode", 0.9, 1.1, GridSampleMode.NEAREST, ValueError), - ] + [("no_min_zoom", None, 1.1, "bilinear", TypeError), ("invalid_mode", 0.9, 1.1, "s", ValueError)] ) def test_invalid_inputs(self, _, min_zoom, max_zoom, mode, raises): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: with self.assertRaises(raises): random_zoom = RandZoom(prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, mode=mode) random_zoom(p(self.imt[0])) def test_auto_expand_3d(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: random_zoom = RandZoom(prob=1.0, min_zoom=[0.8, 0.7], max_zoom=[1.2, 1.3], mode="nearest", keep_size=False) random_zoom.set_random_state(1234) test_data = p(np.random.randint(0, 2, size=[2, 2, 3, 4])) zoomed = random_zoom(test_data) - assert_allclose(random_zoom._zoom, (1.048844, 1.048844, 0.962637), atol=1e-2) - assert_allclose(zoomed.shape, (2, 2, 3, 3)) + assert_allclose(random_zoom._zoom, (1.048844, 1.048844, 0.962637), atol=1e-2, type_test=False) + assert_allclose(zoomed.shape, (2, 2, 3, 3), type_test=False) if __name__ == "__main__": diff --git a/tests/test_rand_zoomd.py b/tests/test_rand_zoomd.py index a22f2f36f1..b2ae40530a 100644 --- a/tests/test_rand_zoomd.py +++ b/tests/test_rand_zoomd.py @@ -16,7 +16,7 @@ from scipy.ndimage import zoom as zoom_scipy from monai.transforms import RandZoomd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion VALID_CASES = [(0.8, 1.2, "nearest", None, False)] @@ -34,25 +34,29 @@ def test_correct_results(self, min_zoom, max_zoom, mode, align_corners, keep_siz align_corners=align_corners, keep_size=keep_size, ) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: random_zoom.set_random_state(1234) - zoomed = random_zoom({key: p(self.imt[0])}) + im = p(self.imt[0]) + zoomed = random_zoom({key: im}) + test_local_inversion(random_zoom, zoomed, {key: im}, key) expected = [ zoom_scipy(channel, zoom=random_zoom.rand_zoom._zoom, mode="nearest", order=0, prefilter=False) for channel in self.imt[0] ] expected = np.stack(expected).astype(np.float32) - assert_allclose(zoomed[key], p(expected), atol=1.0) + assert_allclose(zoomed[key], p(expected), atol=1.0, type_test=False) def test_keep_size(self): key = "img" random_zoom = RandZoomd( keys=key, prob=1.0, min_zoom=0.6, max_zoom=0.7, keep_size=True, padding_mode="constant", constant_values=2 ) - for p in TEST_NDARRAYS: - zoomed = random_zoom({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + zoomed = random_zoom({key: im}) + test_local_inversion(random_zoom, zoomed, {key: im}, key) np.testing.assert_array_equal(zoomed[key].shape, self.imt.shape[1:]) @parameterized.expand( @@ -60,7 +64,7 @@ def test_keep_size(self): ) def test_invalid_inputs(self, _, min_zoom, max_zoom, mode, raises): key = "img" - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: with self.assertRaises(raises): random_zoom = RandZoomd(key, prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, mode=mode) random_zoom({key: p(self.imt[0])}) @@ -69,7 +73,7 @@ def test_auto_expand_3d(self): random_zoom = RandZoomd( keys="img", prob=1.0, min_zoom=[0.8, 0.7], max_zoom=[1.2, 1.3], mode="nearest", keep_size=False ) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: random_zoom.set_random_state(1234) test_data = {"img": p(np.random.randint(0, 2, size=[2, 2, 3, 4]))} zoomed = random_zoom(test_data) diff --git a/tests/test_randtorchvisiond.py b/tests/test_randtorchvisiond.py index 2e96d723ee..4fc1a1e630 100644 --- a/tests/test_randtorchvisiond.py +++ b/tests/test_randtorchvisiond.py @@ -16,7 +16,6 @@ from monai.transforms import Randomizable, RandTorchVisiond from monai.utils import set_determinism -from tests.utils import SkipIfBeforePyTorchVersion TEST_CASE_1 = [ {"keys": "img", "name": "ColorJitter"}, @@ -49,7 +48,6 @@ ] -@SkipIfBeforePyTorchVersion((1, 7)) class TestRandTorchVisiond(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) def test_value(self, input_param, input_data, expected_value): diff --git a/tests/test_remove_repeated_channel.py b/tests/test_remove_repeated_channel.py index 39b42cc4b0..e4b707ce42 100644 --- a/tests/test_remove_repeated_channel.py +++ b/tests/test_remove_repeated_channel.py @@ -11,15 +11,14 @@ import unittest -import numpy as np -import torch from parameterized import parameterized from monai.transforms import RemoveRepeatedChannel +from tests.utils import TEST_NDARRAYS TEST_CASES = [] -for q in (torch.Tensor, np.array): - TEST_CASES.append([{"repeats": 2}, q([[1, 2], [1, 2], [3, 4], [3, 4]]), (2, 2)]) # type: ignore +for q in TEST_NDARRAYS: + TEST_CASES.append([{"repeats": 2}, q([[1, 2], [1, 2], [3, 4], [3, 4]]), (2, 2)]) class TestRemoveRepeatedChannel(unittest.TestCase): diff --git a/tests/test_replace_module.py b/tests/test_replace_module.py new file mode 100644 index 0000000000..4cb4443410 --- /dev/null +++ b/tests/test_replace_module.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. + +import unittest +from typing import Optional, Type + +import torch +from parameterized import parameterized + +from monai.networks.nets import DenseNet121 +from monai.networks.utils import replace_modules, replace_modules_temp +from tests.utils import TEST_DEVICES + +TESTS = [] +for device in TEST_DEVICES: + for match_device in (True, False): + # replace 1 + TESTS.append(("features.denseblock1.denselayer1.layers.relu1", True, match_device, *device)) + # replace 1 (but not strict) + TESTS.append(("features.denseblock1.denselayer1.layers.relu1", False, match_device, *device)) + # replace multiple + TESTS.append(("relu", False, match_device, *device)) + + +class TestReplaceModule(unittest.TestCase): + def setUp(self): + self.net = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) + self.num_relus = self.get_num_modules(torch.nn.ReLU) + self.total = self.get_num_modules() + self.assertGreater(self.num_relus, 0) + + def get_num_modules(self, mod: Optional[Type[torch.nn.Module]] = None) -> int: + m = [m for _, m in self.net.named_modules()] + if mod is not None: + m = [_m for _m in m if isinstance(_m, mod)] + return len(m) + + def check_replaced_modules(self, name, match_device): + # total num modules should remain the same + self.assertEqual(self.total, self.get_num_modules()) + num_relus_mod = self.get_num_modules(torch.nn.ReLU) + num_softmax = self.get_num_modules(torch.nn.Softmax) + # list of returned modules should be as long as number of softmax + self.assertEqual(self.num_relus, num_relus_mod + num_softmax) + if name == "relu": + # at least 2 softmaxes + self.assertGreaterEqual(num_softmax, 2) + else: + # one softmax + self.assertEqual(num_softmax, 1) + if match_device: + self.assertEqual(len(list({i.device for i in self.net.parameters()})), 1) + + @parameterized.expand(TESTS) + def test_replace(self, name, strict_match, match_device, device): + self.net.to(device) + # replace module(s) + replaced = replace_modules(self.net, name, torch.nn.Softmax(), strict_match, match_device) + self.check_replaced_modules(name, match_device) + # number of returned modules should equal number of softmax modules + self.assertEqual(len(replaced), self.get_num_modules(torch.nn.Softmax)) + # all replaced modules should be ReLU + for r in replaced: + self.assertIsInstance(r[1], torch.nn.ReLU) + # if a specfic module was named, check that the name matches exactly + if name == "features.denseblock1.denselayer1.layers.relu1": + self.assertEqual(replaced[0][0], name) + + @parameterized.expand(TESTS) + def test_replace_context_manager(self, name, strict_match, match_device, device): + self.net.to(device) + with replace_modules_temp(self.net, name, torch.nn.Softmax(), strict_match, match_device): + self.check_replaced_modules(name, match_device) + # Check that model was correctly reverted + self.assertEqual(self.get_num_modules(), self.total) + self.assertEqual(self.get_num_modules(torch.nn.ReLU), self.num_relus) + self.assertEqual(self.get_num_modules(torch.nn.Softmax), 0) + + def test_raises(self): + # name doesn't exist in module + with self.assertRaises(AttributeError): + replace_modules(self.net, "non_existent_module", torch.nn.Softmax(), strict_match=True) + with self.assertRaises(AttributeError): + with replace_modules_temp(self.net, "non_existent_module", torch.nn.Softmax(), strict_match=True): + pass + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_resample_to_match.py b/tests/test_resample_to_match.py index e1f6a28998..f1d58e6379 100644 --- a/tests/test_resample_to_match.py +++ b/tests/test_resample_to_match.py @@ -9,9 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import itertools import os +import random +import shutil +import string import tempfile import unittest @@ -27,34 +29,62 @@ TEST_CASES = ["itkreader", "nibabelreader"] +def get_rand_fname(len=10, suffix=".nii.gz"): + letters = string.ascii_letters + out = "".join(random.choice(letters) for _ in range(len)) + out += suffix + return out + + class TestResampleToMatch(unittest.TestCase): - def setUp(self): - self.fnames = [] + @classmethod + def setUpClass(cls): + super(__class__, cls).setUpClass() + cls.fnames = [] + cls.tmpdir = tempfile.mkdtemp() for key in ("0000_t2_tse_tra_4", "0000_ep2d_diff_tra_7"): - fname = os.path.join(os.path.dirname(__file__), "testing_data", f"test_{key}.nii.gz") + fname = os.path.join(cls.tmpdir, f"test_{key}.nii.gz") url = testing_data_config("images", key, "url") hash_type = testing_data_config("images", key, "hash_type") hash_val = testing_data_config("images", key, "hash_val") download_url_or_skip_test(url=url, filepath=fname, hash_type=hash_type, hash_val=hash_val) - self.fnames.append(fname) + cls.fnames.append(fname) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tmpdir) + super(__class__, cls).tearDownClass() @parameterized.expand(itertools.product([NibabelReader, ITKReader], ["monai.data.NibabelWriter", ITKWriter])) def test_correct(self, reader, writer): - with tempfile.TemporaryDirectory() as temp_dir: - loader = Compose([LoadImaged(("im1", "im2"), reader=reader), EnsureChannelFirstd(("im1", "im2"))]) - data = loader({"im1": self.fnames[0], "im2": self.fnames[1]}) - - im_mod, meta = ResampleToMatch()(data["im2"], data["im2_meta_dict"], data["im1_meta_dict"]) - current_dims = copy.deepcopy(meta.get("dim")) - saver = SaveImaged("im3", output_dir=temp_dir, output_postfix="", separate_folder=False, writer=writer) - meta["filename_or_obj"] = "file3.nii.gz" - saver({"im3": im_mod, "im3_meta_dict": meta}) - - saved = nib.load(os.path.join(temp_dir, meta["filename_or_obj"])) - assert_allclose(data["im1"].shape[1:], saved.shape) - assert_allclose(saved.header["dim"][:4], np.array([3, 384, 384, 19])) - if current_dims is not None: - assert_allclose(saved.header["dim"], current_dims) + loader = Compose([LoadImaged(("im1", "im2"), reader=reader), EnsureChannelFirstd(("im1", "im2"))]) + data = loader({"im1": self.fnames[0], "im2": self.fnames[1]}) + + with self.assertRaises(ValueError): + ResampleToMatch(mode=None)(img=data["im2"], img_dst=data["im1"]) + im_mod = ResampleToMatch()(data["im2"], data["im1"]) + saver = SaveImaged( + "im3", output_dir=self.tmpdir, output_postfix="", separate_folder=False, writer=writer, resample=False + ) + im_mod.meta["filename_or_obj"] = get_rand_fname() + saver({"im3": im_mod}) + + saved = nib.load(os.path.join(self.tmpdir, im_mod.meta["filename_or_obj"])) + assert_allclose(data["im1"].shape[1:], saved.shape) + assert_allclose(saved.header["dim"][:4], np.array([3, 384, 384, 19])) + + def test_inverse(self): + loader = Compose([LoadImaged(("im1", "im2")), EnsureChannelFirstd(("im1", "im2"))]) + data = loader({"im1": self.fnames[0], "im2": self.fnames[1]}) + tr = ResampleToMatch() + im_mod = tr(data["im2"], data["im1"]) + self.assertNotEqual(im_mod.shape, data["im2"].shape) + self.assertGreater(((im_mod.affine - data["im2"].affine) ** 2).sum() ** 0.5, 1e-2) + # inverse + im_mod2 = tr.inverse(im_mod) + self.assertEqual(im_mod2.shape, data["im2"].shape) + self.assertLess(((im_mod2.affine - data["im2"].affine) ** 2).sum() ** 0.5, 1e-2) + self.assertEqual(im_mod2.applied_operations, []) if __name__ == "__main__": diff --git a/tests/test_resample_to_matchd.py b/tests/test_resample_to_matchd.py index d9dbeee133..566ef4ada9 100644 --- a/tests/test_resample_to_matchd.py +++ b/tests/test_resample_to_matchd.py @@ -10,6 +10,7 @@ # limitations under the License. import os +import shutil import tempfile import unittest @@ -27,52 +28,51 @@ def update_fname(d): - d["im3_meta_dict"]["filename_or_obj"] = "file3.nii.gz" + d["im3"].meta["filename_or_obj"] = "file3.nii.gz" return d class TestResampleToMatchd(unittest.TestCase): - def setUp(self): - self.fnames = [] + @classmethod + def setUpClass(cls): + super(__class__, cls).setUpClass() + cls.fnames = [] + cls.tmpdir = tempfile.mkdtemp() for key in ("0000_t2_tse_tra_4", "0000_ep2d_diff_tra_7"): - fname = os.path.join(os.path.dirname(__file__), "testing_data", f"test_{key}.nii.gz") + fname = os.path.join(cls.tmpdir, f"test_{key}.nii.gz") url = testing_data_config("images", key, "url") hash_type = testing_data_config("images", key, "hash_type") hash_val = testing_data_config("images", key, "hash_val") download_url_or_skip_test(url=url, filepath=fname, hash_type=hash_type, hash_val=hash_val) - self.fnames.append(fname) + cls.fnames.append(fname) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tmpdir) + super(__class__, cls).tearDownClass() def test_correct(self): - with tempfile.TemporaryDirectory() as temp_dir: - transforms = Compose( - [ - LoadImaged(("im1", "im2")), - EnsureChannelFirstd(("im1", "im2")), - CopyItemsd(("im2", "im2_meta_dict"), names=("im3", "im3_meta_dict")), - ResampleToMatchd("im3", "im1_meta_dict"), - Lambda(update_fname), - SaveImaged("im3", output_dir=temp_dir, output_postfix="", separate_folder=False), - ] - ) - data = transforms({"im1": self.fnames[0], "im2": self.fnames[1]}) - # check that output sizes match - assert_allclose(data["im1"].shape, data["im3"].shape) - # and that the meta data has been updated accordingly - assert_allclose(data["im3"].shape[1:], data["im3_meta_dict"]["spatial_shape"], type_test=False) - assert_allclose(data["im3_meta_dict"]["affine"], data["im1_meta_dict"]["affine"]) - # check we're different from the original - self.assertTrue(any(i != j for i, j in zip(data["im3"].shape, data["im2"].shape))) - self.assertTrue( - any( - i != j - for i, j in zip( - data["im3_meta_dict"]["affine"].flatten(), data["im2_meta_dict"]["affine"].flatten() - ) - ) - ) - # test the inverse - data = Invertd("im3", transforms, "im3")(data) - assert_allclose(data["im2"].shape, data["im3"].shape) + transforms = Compose( + [ + LoadImaged(("im1", "im2")), + EnsureChannelFirstd(("im1", "im2")), + CopyItemsd(("im2"), names=("im3")), + ResampleToMatchd("im3", "im1"), + Lambda(update_fname), + SaveImaged("im3", output_dir=self.tmpdir, output_postfix="", separate_folder=False, resample=False), + ] + ) + data = transforms({"im1": self.fnames[0], "im2": self.fnames[1]}) + # check that output sizes match + assert_allclose(data["im1"].shape, data["im3"].shape) + # and that the meta data has been updated accordingly + assert_allclose(data["im3"].affine, data["im1"].affine) + # check we're different from the original + self.assertTrue(any(i != j for i, j in zip(data["im3"].shape, data["im2"].shape))) + self.assertTrue(any(i != j for i, j in zip(data["im3"].affine.flatten(), data["im2"].affine.flatten()))) + # test the inverse + data = Invertd("im3", transforms)(data) + assert_allclose(data["im2"].shape, data["im3"].shape) if __name__ == "__main__": diff --git a/tests/test_resampler.py b/tests/test_resampler.py index 7dfb86a7a9..5c8ef24c0e 100644 --- a/tests/test_resampler.py +++ b/tests/test_resampler.py @@ -17,11 +17,11 @@ from monai.transforms import Resample from monai.transforms.utils import create_grid -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: + for q in TEST_NDARRAYS_ALL: for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]: TESTS.append( [ @@ -156,7 +156,7 @@ def test_resample(self, input_param, input_data, expected_val): result = g(**input_data) if "device" in input_data: self.assertEqual(result.device, input_data["device"]) - assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4) + assert_allclose(result, expected_val, rtol=1e-4, atol=1e-4, type_test=False) if __name__ == "__main__": diff --git a/tests/test_resize.py b/tests/test_resize.py index 06246b2358..8927b5dba5 100644 --- a/tests/test_resize.py +++ b/tests/test_resize.py @@ -13,10 +13,12 @@ import numpy as np import skimage.transform +import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Resize -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, is_tf32_env, pytorch_after TEST_CASE_0 = [{"spatial_size": 15}, (6, 10, 15)] @@ -24,6 +26,12 @@ TEST_CASE_2 = [{"spatial_size": 6, "mode": "trilinear", "align_corners": True}, (2, 4, 6)] +TEST_CASE_3 = [{"spatial_size": 15, "anti_aliasing": True}, (6, 10, 15)] + +TEST_CASE_4 = [{"spatial_size": 6, "anti_aliasing": True, "anti_aliasing_sigma": 2.0}, (2, 4, 6)] + +diff_t = 0.3 if is_tf32_env() else 0.2 + class TestResize(NumpyImageTestCase2D): def test_invalid_inputs(self): @@ -36,34 +44,68 @@ def test_invalid_inputs(self): resize(self.imt[0]) @parameterized.expand( - [((32, -1), "area"), ((32, 32), "area"), ((32, 32, 32), "trilinear"), ((256, 256), "bilinear")] + [ + ((32, -1), "area", True), + ((32, 32), "area", False), + ((32, 32, 32), "trilinear", True), + ((256, 256), "bilinear", False), + ((256, 256), "nearest-exact" if pytorch_after(1, 11) else "nearest", False), + ((128, 64), "area", True), # already in a good shape + ] ) - def test_correct_results(self, spatial_size, mode): - resize = Resize(spatial_size, mode=mode) + def test_correct_results(self, spatial_size, mode, anti_aliasing): + """resize 'spatial_size' and 'mode'""" + resize = Resize(spatial_size, mode=mode, anti_aliasing=anti_aliasing) _order = 0 if mode.endswith("linear"): _order = 1 if spatial_size == (32, -1): spatial_size = (32, 64) + expected = [ skimage.transform.resize( - channel, spatial_size, order=_order, clip=False, preserve_range=False, anti_aliasing=False + channel, spatial_size, order=_order, clip=False, preserve_range=False, anti_aliasing=anti_aliasing ) for channel in self.imt[0] ] expected = np.stack(expected).astype(np.float32) - for p in TEST_NDARRAYS: - out = resize(p(self.imt[0])) - assert_allclose(out, expected, type_test=False, atol=0.9) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + out = resize(im) + if isinstance(im, MetaTensor): + if not out.applied_operations: + return # skipped because good shape + im_inv = resize.inverse(out) + self.assertTrue(not im_inv.applied_operations) + assert_allclose(im_inv.shape, im.shape) + assert_allclose(im_inv.affine, im.affine, atol=1e-3, rtol=1e-3) + if not anti_aliasing: + assert_allclose(out, expected, type_test=False, atol=0.9) + return + # skimage uses reflect padding for anti-aliasing filter. + # Our implementation reuses GaussianSmooth() as anti-aliasing filter, which uses zero padding instead. + # Thus their results near the image boundary will be different. + if isinstance(out, torch.Tensor): + out = out.cpu().detach().numpy() + good = np.sum(np.isclose(expected, out, atol=0.9)) + self.assertLessEqual( + np.abs(good - expected.size) / float(expected.size), diff_t, f"at most {diff_t} percent mismatch " + ) - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2]) + @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) def test_longest_shape(self, input_param, expected_shape): input_data = np.random.randint(0, 2, size=[3, 4, 7, 10]) input_param["size_mode"] = "longest" result = Resize(**input_param)(input_data) np.testing.assert_allclose(result.shape[1:], expected_shape) + set_track_meta(False) + result = Resize(**input_param)(input_data) + self.assertNotIsInstance(result, MetaTensor) + np.testing.assert_allclose(result.shape[1:], expected_shape) + set_track_meta(True) + def test_longest_infinite_decimals(self): resize = Resize(spatial_size=1008, size_mode="longest", mode="bilinear", align_corners=False) ret = resize(np.random.randint(0, 2, size=[1, 2544, 3032])) diff --git a/tests/test_resize_with_pad_or_crop.py b/tests/test_resize_with_pad_or_crop.py index f81e1d4b08..4e097cd3d4 100644 --- a/tests/test_resize_with_pad_or_crop.py +++ b/tests/test_resize_with_pad_or_crop.py @@ -15,8 +15,9 @@ import torch from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms import ResizeWithPadOrCrop -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL, pytorch_after TEST_CASES = [ [{"spatial_size": [15, 8, 8], "mode": "constant"}, (3, 8, 8, 4), (3, 15, 8, 8)], @@ -26,24 +27,38 @@ (3, 15, 4, 8), ], [{"spatial_size": [15, 4, -1], "mode": "constant"}, (3, 8, 8, 4), (3, 15, 4, 4)], - [{"spatial_size": [15, 4, -1], "mode": "reflect"}, (3, 8, 8, 4), (3, 15, 4, 4)], - [{"spatial_size": [-1, -1, -1], "mode": "reflect"}, (3, 8, 8, 4), (3, 8, 8, 4)], + [ + {"spatial_size": [15, 4, -1], "mode": "reflect" if pytorch_after(1, 11) else "constant"}, + (3, 8, 8, 4), + (3, 15, 4, 4), + ], + [ + {"spatial_size": [-1, -1, -1], "mode": "reflect" if pytorch_after(1, 11) else "constant"}, + (3, 8, 8, 4), + (3, 8, 8, 4), + ], ] class TestResizeWithPadOrCrop(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_pad_shape(self, input_param, input_shape, expected_shape): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: if isinstance(p(0), torch.Tensor) and ( "constant_values" in input_param or input_param["mode"] == "reflect" ): continue - paddcroper = ResizeWithPadOrCrop(**input_param) - result = paddcroper(p(np.zeros(input_shape))) + padcropper = ResizeWithPadOrCrop(**input_param) + result = padcropper(p(np.zeros(input_shape))) np.testing.assert_allclose(result.shape, expected_shape) - result = paddcroper(p(np.zeros(input_shape)), mode="constant") + result = padcropper(p(np.zeros(input_shape)), mode="constant") np.testing.assert_allclose(result.shape, expected_shape) + self.assertIsInstance(result, MetaTensor) + self.assertEqual(len(result.applied_operations), 1) + inv = padcropper.inverse(result) + self.assertTupleEqual(inv.shape, input_shape) + self.assertIsInstance(inv, MetaTensor) + self.assertEqual(inv.applied_operations, []) if __name__ == "__main__": diff --git a/tests/test_resize_with_pad_or_cropd.py b/tests/test_resize_with_pad_or_cropd.py index 28993a2bf4..eb4e5f09cc 100644 --- a/tests/test_resize_with_pad_or_cropd.py +++ b/tests/test_resize_with_pad_or_cropd.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.transforms import ResizeWithPadOrCropd -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS_ALL, pytorch_after TEST_CASES = [ [{"keys": "img", "spatial_size": [15, 8, 8], "mode": "constant"}, {"img": np.zeros((3, 8, 8, 4))}, (3, 15, 8, 8)], @@ -26,23 +26,34 @@ (3, 15, 4, 8), ], [{"keys": "img", "spatial_size": [15, 4, -1], "mode": "constant"}, {"img": np.zeros((3, 8, 8, 4))}, (3, 15, 4, 4)], - [{"keys": "img", "spatial_size": [15, 4, -1], "mode": "reflect"}, {"img": np.zeros((3, 8, 8, 4))}, (3, 15, 4, 4)], - [{"keys": "img", "spatial_size": [-1, -1, -1], "mode": "reflect"}, {"img": np.zeros((3, 8, 8, 4))}, (3, 8, 8, 4)], + [ + {"keys": "img", "spatial_size": [15, 4, -1], "mode": "reflect" if pytorch_after(1, 11) else "constant"}, + {"img": np.zeros((3, 8, 8, 4))}, + (3, 15, 4, 4), + ], + [ + {"keys": "img", "spatial_size": [-1, -1, -1], "mode": "reflect" if pytorch_after(1, 11) else "constant"}, + {"img": np.zeros((3, 8, 8, 4))}, + (3, 8, 8, 4), + ], ] class TestResizeWithPadOrCropd(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_pad_shape(self, input_param, input_data, expected_val): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: if isinstance(p(0), torch.Tensor) and ( "constant_values" in input_param or input_param["mode"] == "reflect" ): continue - paddcroper = ResizeWithPadOrCropd(**input_param) + padcropper = ResizeWithPadOrCropd(**input_param) input_data["img"] = p(input_data["img"]) - result = paddcroper(input_data) + result = padcropper(input_data) np.testing.assert_allclose(result["img"].shape, expected_val) + inv = padcropper.inverse(result) + for k in input_data: + self.assertTupleEqual(inv[k].shape, input_data[k].shape) if __name__ == "__main__": diff --git a/tests/test_resized.py b/tests/test_resized.py index d7374ea930..b8db666357 100644 --- a/tests/test_resized.py +++ b/tests/test_resized.py @@ -15,8 +15,9 @@ import skimage.transform from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Resized -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion TEST_CASE_0 = [{"keys": "img", "spatial_size": 15}, (6, 10, 15)] @@ -56,9 +57,11 @@ def test_correct_results(self, spatial_size, mode): ] expected = np.stack(expected).astype(np.float32) - for p in TEST_NDARRAYS: - out = resize({"img": p(self.imt[0])})["img"] - assert_allclose(out, expected, type_test=False, atol=0.9) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + out = resize({"img": im}) + test_local_inversion(resize, out, {"img": im}, "img") + assert_allclose(out["img"], expected, type_test=False, atol=0.9) @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) def test_longest_shape(self, input_param, expected_shape): @@ -71,6 +74,11 @@ def test_longest_shape(self, input_param, expected_shape): result = rescaler(input_data) for k in rescaler.keys: np.testing.assert_allclose(result[k].shape[1:], expected_shape) + set_track_meta(False) + result = Resized(**input_param)(input_data) + self.assertNotIsInstance(result["img"], MetaTensor) + np.testing.assert_allclose(result["img"].shape[1:], expected_shape) + set_track_meta(True) if __name__ == "__main__": diff --git a/tests/test_retinanet.py b/tests/test_retinanet.py new file mode 100644 index 0000000000..3c136a4cf2 --- /dev/null +++ b/tests/test_retinanet.py @@ -0,0 +1,175 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.networks.retinanet_network import RetinaNet, resnet_fpn_feature_extractor +from monai.networks import eval_mode +from monai.networks.nets import resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200 +from monai.utils import ensure_tuple, optional_import +from tests.utils import SkipIfBeforePyTorchVersion, test_script_save + +_, has_torchvision = optional_import("torchvision") + + +device = "cuda" if torch.cuda.is_available() else "cpu" +num_anchors = 7 + +TEST_CASE_1 = [ # 3D, batch 3, 2 input channel + { + "pretrained": False, + "spatial_dims": 3, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": 7, + "conv1_t_stride": (2, 2, 2), + }, + (3, 2, 32, 64, 48), +] + +TEST_CASE_2 = [ # 2D, batch 2, 1 input channel + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "conv1_t_size": [7, 7], + "conv1_t_stride": [2, 2], + }, + (2, 1, 32, 64), +] + +TEST_CASE_2_A = [ # 2D, batch 2, 1 input channel, shortcut type A + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "shortcut_type": "A", + "conv1_t_size": (7, 7), + "conv1_t_stride": 2, + }, + (2, 1, 32, 64), +] + +TEST_CASE_3 = [ # 1D, batch 1, 2 input channels + { + "pretrained": False, + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), +] + +TEST_CASE_3_A = [ # 1D, batch 1, 2 input channels + {"pretrained": False, "spatial_dims": 1, "n_input_channels": 2, "num_classes": 3, "shortcut_type": "A"}, + (1, 2, 32), +] + +TEST_CASE_4 = [ # 2D, batch 2, 1 input channel + {"pretrained": False, "spatial_dims": 2, "n_input_channels": 1, "num_classes": 3, "feed_forward": False}, + (2, 1, 32, 64), +] + +TEST_CASES = [] +for case in [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_2_A, TEST_CASE_3_A]: + for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200]: + TEST_CASES.append([model, *case]) + +TEST_CASES_TS = [] +for case in [TEST_CASE_1]: + for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200]: + TEST_CASES_TS.append([model, *case]) + + +@SkipIfBeforePyTorchVersion((1, 9)) +@unittest.skipUnless(has_torchvision, "Requires torchvision") +class TestRetinaNet(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_retina_shape(self, model, input_param, input_shape): + backbone = model(**input_param) + feature_extractor = resnet_fpn_feature_extractor( + backbone=backbone, + spatial_dims=input_param["spatial_dims"], + pretrained_backbone=input_param["pretrained"], + trainable_backbone_layers=None, + returned_layers=[1, 2], + ) + net = RetinaNet( + spatial_dims=input_param["spatial_dims"], + num_classes=input_param["num_classes"], + num_anchors=num_anchors, + feature_extractor=feature_extractor, + size_divisible=32, + ).to(device) + + with eval_mode(net): + result = net.forward(torch.randn(input_shape).to(device)) + + base_stride = ensure_tuple(input_param["conv1_t_stride"])[0] if "conv1_t_stride" in input_param else 1 + expected_cls_channel = input_param["num_classes"] * num_anchors + expected_cls_shape = tuple( + (input_shape[0], expected_cls_channel) + + tuple(input_shape[2 + a] // s // base_stride for a in range(input_param["spatial_dims"])) + for s in [2, 4, 8] + ) + expected_box_channel = 2 * input_param["spatial_dims"] * num_anchors + expected_box_shape = tuple( + (input_shape[0], expected_box_channel) + + tuple(input_shape[2 + a] // s // base_stride for a in range(input_param["spatial_dims"])) + for s in [2, 4, 8] + ) + + self.assertEqual(tuple(cc.shape for cc in result[net.cls_key]), expected_cls_shape) + self.assertEqual(tuple(cc.shape for cc in result[net.box_reg_key]), expected_box_shape) + + @parameterized.expand(TEST_CASES_TS) + def test_script(self, model, input_param, input_shape): + try: + idx = int(self.id().split("test_script_")[-1]) + except BaseException: + idx = 0 + idx %= 3 + # test whether support torchscript + data = torch.randn(input_shape) + backbone = model(**input_param) + if idx == 0: + test_script_save(backbone, data) + return + feature_extractor = resnet_fpn_feature_extractor( + backbone=backbone, + spatial_dims=input_param["spatial_dims"], + pretrained_backbone=input_param["pretrained"], + trainable_backbone_layers=None, + returned_layers=[1, 2], + ) + if idx == 1: + test_script_save(feature_extractor, data) + return + net = RetinaNet( + spatial_dims=input_param["spatial_dims"], + num_classes=input_param["num_classes"], + num_anchors=num_anchors, + feature_extractor=feature_extractor, + size_divisible=32, + ) + if idx == 2: + test_script_save(net, data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_retinanet_detector.py b/tests/test_retinanet_detector.py new file mode 100644 index 0000000000..99a70fb5fa --- /dev/null +++ b/tests/test_retinanet_detector.py @@ -0,0 +1,200 @@ +# 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. + +import random +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.networks.retinanet_detector import RetinaNetDetector, retinanet_resnet50_fpn_detector +from monai.apps.detection.utils.anchor_utils import AnchorGeneratorWithAnchorShape +from monai.networks import eval_mode, train_mode +from monai.utils import optional_import +from tests.utils import SkipIfBeforePyTorchVersion, test_script_save + +_, has_torchvision = optional_import("torchvision") + + +num_anchors = 7 + +TEST_CASE_1 = [ # 3D, batch 3, 2 input channel + { + "pretrained": False, + "spatial_dims": 3, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": 7, + "conv1_t_stride": (2, 2, 2), + }, + (3, 2, 32, 64, 48), +] + +TEST_CASE_2 = [ # 2D, batch 2, 1 input channel + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "conv1_t_size": [7, 7], + "conv1_t_stride": [2, 2], + }, + (2, 1, 32, 64), +] + +TEST_CASE_2_A = [ # 2D, batch 2, 1 input channel, shortcut type A + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "shortcut_type": "A", + "conv1_t_size": (7, 7), + "conv1_t_stride": 2, + }, + (2, 1, 32, 64), +] + +TEST_CASE_3 = [ # 1D, batch 1, 2 input channels + { + "pretrained": False, + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), +] + +TEST_CASE_3_A = [ # 1D, batch 1, 2 input channels + {"pretrained": False, "spatial_dims": 1, "n_input_channels": 2, "num_classes": 3, "shortcut_type": "A"}, + (1, 2, 32), +] + +TEST_CASE_4 = [ # 2D, batch 2, 1 input channel + {"pretrained": False, "spatial_dims": 2, "n_input_channels": 1, "num_classes": 3, "feed_forward": False}, + (2, 1, 32, 64), +] + +TEST_CASES = [] +TEST_CASES = [TEST_CASE_1, TEST_CASE_2, TEST_CASE_2_A] + +TEST_CASES_TS = [TEST_CASE_1] + + +class NaiveNetwork(torch.nn.Module): + def __init__(self, spatial_dims, num_classes, **kwargs): + super().__init__() + self.spatial_dims = spatial_dims + self.num_classes = num_classes + self.num_anchors = 1 + self.cls_key = "cls" + self.box_reg_key = "box_reg" + self.size_divisible = 1 + + def forward(self, images): + out_cls_shape = (images.shape[0], self.num_classes * self.num_anchors) + images.shape[-self.spatial_dims :] + out_box_reg_shape = (images.shape[0], 2 * self.spatial_dims * self.num_anchors) + images.shape[ + -self.spatial_dims : + ] + return {self.cls_key: [torch.randn(out_cls_shape)], self.box_reg_key: [torch.randn(out_box_reg_shape)]} + + +@SkipIfBeforePyTorchVersion((1, 11)) +@unittest.skipUnless(has_torchvision, "Requires torchvision") +class TestRetinaNetDetector(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_retina_detector_resnet_backbone_shape(self, input_param, input_shape): + returned_layers = [1] + anchor_generator = AnchorGeneratorWithAnchorShape( + feature_map_scales=(1, 2), base_anchor_shapes=((8,) * input_param["spatial_dims"],) + ) + detector = retinanet_resnet50_fpn_detector( + **input_param, anchor_generator=anchor_generator, returned_layers=returned_layers + ) + + with eval_mode(detector): + input_data = torch.randn(input_shape) + result = detector.forward(input_data) + assert len(result) == len(result) + + input_data = [torch.randn(input_shape[1:]) for _ in range(random.randint(1, 9))] + result = detector.forward(input_data) + assert len(result) == len(result) + + detector.set_atss_matcher() + detector.set_hard_negative_sampler(10, 0.5) + gt_box_start = torch.randint(2, (3, input_param["spatial_dims"])).to(torch.float16) + gt_box_end = gt_box_start + torch.randint(1, 10, (3, input_param["spatial_dims"])) + one_target = { + "boxes": torch.cat((gt_box_start, gt_box_end), dim=1), + "labels": torch.randint(input_param["num_classes"], (3,)), + } + with train_mode(detector): + input_data = torch.randn(input_shape) + targets = [one_target] * len(input_data) + result = detector.forward(input_data, targets) + + input_data = [torch.randn(input_shape[1:]) for _ in range(random.randint(1, 9))] + targets = [one_target] * len(input_data) + result = detector.forward(input_data, targets) + + @parameterized.expand(TEST_CASES) + def test_naive_retina_detector_shape(self, input_param, input_shape): + anchor_generator = AnchorGeneratorWithAnchorShape( + feature_map_scales=(1,), base_anchor_shapes=((8,) * input_param["spatial_dims"],) + ) + detector = RetinaNetDetector(network=NaiveNetwork(**input_param), anchor_generator=anchor_generator) + + with eval_mode(detector): + input_data = torch.randn(input_shape) + result = detector.forward(input_data) + assert len(result) == len(result) + + input_data = [torch.randn(input_shape[1:]) for _ in range(random.randint(1, 9))] + result = detector.forward(input_data) + assert len(result) == len(result) + + detector.set_atss_matcher() + detector.set_hard_negative_sampler(10, 0.5) + gt_box_start = torch.randint(2, (3, input_param["spatial_dims"])).to(torch.float16) + gt_box_end = gt_box_start + torch.randint(1, 10, (3, input_param["spatial_dims"])) + one_target = { + "boxes": torch.cat((gt_box_start, gt_box_end), dim=1), + "labels": torch.randint(input_param["num_classes"], (3,)), + } + with train_mode(detector): + input_data = torch.randn(input_shape) + targets = [one_target] * len(input_data) + result = detector.forward(input_data, targets) + + input_data = [torch.randn(input_shape[1:]) for _ in range(random.randint(1, 9))] + targets = [one_target] * len(input_data) + result = detector.forward(input_data, targets) + + @parameterized.expand(TEST_CASES_TS) + def test_script(self, input_param, input_shape): + # test whether support torchscript + returned_layers = [1] + anchor_generator = AnchorGeneratorWithAnchorShape( + feature_map_scales=(1, 2), base_anchor_shapes=((8,) * input_param["spatial_dims"],) + ) + detector = retinanet_resnet50_fpn_detector( + **input_param, anchor_generator=anchor_generator, returned_layers=returned_layers + ) + with eval_mode(detector): + input_data = torch.randn(input_shape) + test_script_save(detector.network, input_data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_retinanet_predict_utils.py b/tests/test_retinanet_predict_utils.py new file mode 100644 index 0000000000..5157691696 --- /dev/null +++ b/tests/test_retinanet_predict_utils.py @@ -0,0 +1,149 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.apps.detection.utils.predict_utils import ensure_dict_value_to_list_, predict_with_inferer +from monai.inferers import SlidingWindowInferer + +TEST_CASE_1 = [ # 3D, batch 3, 2 input channel + { + "pretrained": False, + "spatial_dims": 3, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": 7, + "conv1_t_stride": (2, 2, 2), + }, + (3, 2, 32, 64, 48), +] + +TEST_CASE_2 = [ # 2D, batch 2, 1 input channel + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "conv1_t_size": [7, 7], + "conv1_t_stride": [2, 2], + }, + (2, 1, 32, 64), +] + +TEST_CASE_2_A = [ # 2D, batch 2, 1 input channel, shortcut type A + { + "pretrained": False, + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 3, + "shortcut_type": "A", + "conv1_t_size": (7, 7), + "conv1_t_stride": 2, + }, + (2, 1, 32, 64), +] + +TEST_CASE_3 = [ # 1D, batch 1, 2 input channels + { + "pretrained": False, + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), +] + +TEST_CASE_3_A = [ # 1D, batch 1, 2 input channels + {"pretrained": False, "spatial_dims": 1, "n_input_channels": 2, "num_classes": 3, "shortcut_type": "A"}, + (1, 2, 32), +] + +TEST_CASE_4 = [ # 2D, batch 2, 1 input channel + {"pretrained": False, "spatial_dims": 2, "n_input_channels": 1, "num_classes": 3, "feed_forward": False}, + (2, 1, 32, 64), +] + +TEST_CASES = [] +TEST_CASES = [TEST_CASE_1, TEST_CASE_2, TEST_CASE_2_A] + +TEST_CASES_TS = [TEST_CASE_1] + + +class NaiveNetwork(torch.nn.Module): + def __init__(self, spatial_dims, num_classes, **kwargs): + super().__init__() + self.spatial_dims = spatial_dims + self.num_classes = num_classes + self.num_anchors = 1 + self.cls_key = "cls" + self.box_reg_key = "box_reg" + self.size_divisible = 1 + + def forward(self, images): + out_cls_shape = (images.shape[0], self.num_classes * self.num_anchors) + images.shape[-self.spatial_dims :] + out_box_reg_shape = (images.shape[0], 2 * self.spatial_dims * self.num_anchors) + images.shape[ + -self.spatial_dims : + ] + return {self.cls_key: torch.randn(out_cls_shape), self.box_reg_key: [torch.randn(out_box_reg_shape)]} + + +class NaiveNetwork2(torch.nn.Module): + def __init__(self, spatial_dims, num_classes, **kwargs): + super().__init__() + self.spatial_dims = spatial_dims + self.num_classes = num_classes + self.num_anchors = 1 + self.cls_key = "cls" + self.box_reg_key = "box_reg" + self.size_divisible = 1 + + def forward(self, images): + out_cls_shape = (images.shape[0], self.num_classes * self.num_anchors) + images.shape[-self.spatial_dims :] + out_box_reg_shape = (images.shape[0], 2 * self.spatial_dims * self.num_anchors) + images.shape[ + -self.spatial_dims : + ] + return {self.cls_key: [torch.randn(out_cls_shape)] * 2, self.box_reg_key: [torch.randn(out_box_reg_shape)] * 2} + + +class TestPredictor(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_naive_predictor(self, input_param, input_shape): + net = NaiveNetwork(**input_param) + net2 = NaiveNetwork2(**input_param) + inferer = SlidingWindowInferer(roi_size=16, overlap=0.25, cache_roi_weight_map=True) + network_output_keys = ["cls", "box_reg"] + + input_data = torch.randn(input_shape) + + result = predict_with_inferer(input_data, net, network_output_keys, inferer=inferer) + self.assertTrue(len(result["cls"]) == 1) + + result = net(input_data) + self.assertTrue(len(result["cls"]) == input_data.shape[0]) + ensure_dict_value_to_list_(result) + self.assertTrue(len(result["cls"]) == 1) + + result = predict_with_inferer(input_data, net2, network_output_keys, inferer=inferer) + self.assertTrue(len(result["cls"]) == 2) + + result = net2(input_data) + self.assertTrue(len(result["cls"]) == 2) + ensure_dict_value_to_list_(result) + self.assertTrue(len(result["cls"]) == 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_rotate.py b/tests/test_rotate.py index 01842f6d73..d039738b21 100644 --- a/tests/test_rotate.py +++ b/tests/test_rotate.py @@ -17,11 +17,12 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import Rotate -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D, test_local_inversion TEST_CASES_2D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_2D.append((p, np.pi / 6, False, "bilinear", "border", False)) TEST_CASES_2D.append((p, np.pi / 4, True, "bilinear", "border", False)) TEST_CASES_2D.append((p, -np.pi / 4.5, True, "nearest", "reflection", False)) @@ -29,7 +30,7 @@ TEST_CASES_2D.append((p, -np.pi / 2, False, "bilinear", "zeros", True)) TEST_CASES_3D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_3D.append((p, -np.pi / 2, True, "nearest", "border", False)) TEST_CASES_3D.append((p, np.pi / 4, True, "bilinear", "border", False)) TEST_CASES_3D.append((p, -np.pi / 4.5, True, "nearest", "reflection", False)) @@ -37,7 +38,7 @@ TEST_CASES_3D.append((p, -np.pi / 2, False, "bilinear", "zeros", False)) TEST_CASES_SHAPE_3D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_SHAPE_3D.append((p, [-np.pi / 2, 1.0, 2.0], "nearest", "border", False)) TEST_CASES_SHAPE_3D.append((p, [np.pi / 4, 0, 0], "bilinear", "border", False)) TEST_CASES_SHAPE_3D.append((p, [-np.pi / 4.5, -20, 20], "nearest", "reflection", False)) @@ -101,11 +102,18 @@ def test_correct_results(self, im_type, angle, keep_size, mode, padding_mode, al @parameterized.expand(TEST_CASES_SHAPE_3D) def test_correct_shape(self, im_type, angle, mode, padding_mode, align_corners): rotate_fn = Rotate(angle, True, align_corners=align_corners, dtype=np.float64) - rotated = rotate_fn(im_type(self.imt[0]), mode=mode, padding_mode=padding_mode) + im = im_type(self.imt[0]) + set_track_meta(False) + rotated = rotate_fn(im, mode=mode, padding_mode=padding_mode) + self.assertNotIsInstance(rotated, MetaTensor) np.testing.assert_allclose(self.imt[0].shape, rotated.shape) + set_track_meta(True) + rotated = rotate_fn(im, mode=mode, padding_mode=padding_mode) + np.testing.assert_allclose(self.imt[0].shape, rotated.shape) + test_local_inversion(rotate_fn, rotated, im) def test_ill_case(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: rotate_fn = Rotate(10, True) with self.assertRaises(ValueError): # wrong shape rotate_fn(p(self.imt)) diff --git a/tests/test_rotate90.py b/tests/test_rotate90.py index 9865120688..69414430c2 100644 --- a/tests/test_rotate90.py +++ b/tests/test_rotate90.py @@ -13,42 +13,105 @@ import numpy as np +from monai.data import MetaTensor, set_track_meta from monai.transforms import Rotate90 -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import ( + TEST_NDARRAYS_ALL, + NumpyImageTestCase2D, + NumpyImageTestCase3D, + assert_allclose, + test_local_inversion, +) class TestRotate90(NumpyImageTestCase2D): def test_rotate90_default(self): rotate = Rotate90() - for p in TEST_NDARRAYS: - rotated = rotate(p(self.imt[0])) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + set_track_meta(True) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + set_track_meta(False) + rotated = rotate(im) + self.assertNotIsInstance(rotated, MetaTensor) + set_track_meta(True) def test_k(self): rotate = Rotate90(k=2) - for p in TEST_NDARRAYS: - rotated = rotate(p(self.imt[0])) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") def test_spatial_axes(self): rotate = Rotate90(spatial_axes=(0, -1)) - for p in TEST_NDARRAYS: - rotated = rotate(p(self.imt[0])) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 1, (0, -1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") def test_prob_k_spatial_axes(self): rotate = Rotate90(k=2, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: - rotated = rotate(p(self.imt[0])) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) + expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] + expected = np.stack(expected) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + + +class TestRotate903d(NumpyImageTestCase3D): + def test_rotate90_default(self): + rotate = Rotate90() + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) + expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] + expected = np.stack(expected) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + + def test_k(self): + rotate = Rotate90(k=2) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) + expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] + expected = np.stack(expected) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + + def test_spatial_axes(self): + rotate = Rotate90(spatial_axes=(0, -1)) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) + expected = [np.rot90(channel, 1, (0, -1)) for channel in self.imt[0]] + expected = np.stack(expected) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") + + def test_prob_k_spatial_axes(self): + rotate = Rotate90(k=2, spatial_axes=(0, 1)) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate(im) + test_local_inversion(rotate, rotated, im) expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8) + assert_allclose(rotated, p(expected), rtol=1.0e-5, atol=1.0e-8, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_rotate90d.py b/tests/test_rotate90d.py index ef4bad9419..f88e8937e8 100644 --- a/tests/test_rotate90d.py +++ b/tests/test_rotate90d.py @@ -13,46 +13,60 @@ import numpy as np +from monai.data import MetaTensor, set_track_meta from monai.transforms import Rotate90d -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion class TestRotate90d(NumpyImageTestCase2D): def test_rotate90_default(self): key = "test" rotate = Rotate90d(keys=key) - for p in TEST_NDARRAYS: - rotated = rotate({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + set_track_meta(True) + rotated = rotate({key: im}) + test_local_inversion(rotate, rotated, {key: im}, key) expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") + set_track_meta(False) + rotated = rotate({key: im}) + self.assertNotIsInstance(rotated[key], MetaTensor) + set_track_meta(True) def test_k(self): key = None rotate = Rotate90d(keys=key, k=2) - for p in TEST_NDARRAYS: - rotated = rotate({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate({key: im}) + test_local_inversion(rotate, rotated, {key: im}, key) expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") def test_spatial_axes(self): key = "test" rotate = Rotate90d(keys=key, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: - rotated = rotate({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate({key: im}) + test_local_inversion(rotate, rotated, {key: im}, key) expected = [np.rot90(channel, 1, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") def test_prob_k_spatial_axes(self): key = "test" rotate = Rotate90d(keys=key, k=2, spatial_axes=(0, 1)) - for p in TEST_NDARRAYS: - rotated = rotate({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + rotated = rotate({key: im}) + test_local_inversion(rotate, rotated, {key: im}, key) expected = [np.rot90(channel, 2, (0, 1)) for channel in self.imt[0]] expected = np.stack(expected) - assert_allclose(rotated[key], p(expected)) + assert_allclose(rotated[key], p(expected), type_test="tensor") def test_no_key(self): key = "unknown" diff --git a/tests/test_rotated.py b/tests/test_rotated.py index 43b5a68f61..48b2e8a3c7 100644 --- a/tests/test_rotated.py +++ b/tests/test_rotated.py @@ -17,11 +17,12 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor from monai.transforms import Rotated -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, NumpyImageTestCase3D, test_local_inversion TEST_CASES_2D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_2D.append((p, -np.pi / 6, False, "bilinear", "border", False)) TEST_CASES_2D.append((p, -np.pi / 4, True, "bilinear", "border", False)) TEST_CASES_2D.append((p, np.pi / 4.5, True, "nearest", "reflection", False)) @@ -29,7 +30,7 @@ TEST_CASES_2D.append((p, np.pi / 2, False, "bilinear", "zeros", True)) TEST_CASES_3D: List[Tuple] = [] -for p in TEST_NDARRAYS: +for p in TEST_NDARRAYS_ALL: TEST_CASES_3D.append((p, -np.pi / 6, False, "bilinear", "border", False)) TEST_CASES_3D.append((p, -np.pi / 4, True, "bilinear", "border", False)) TEST_CASES_3D.append((p, np.pi / 4.5, True, "nearest", "reflection", False)) @@ -43,7 +44,8 @@ def test_correct_results(self, im_type, angle, keep_size, mode, padding_mode, al rotate_fn = Rotated( ("img", "seg"), angle, keep_size, (mode, "nearest"), padding_mode, align_corners, dtype=np.float64 ) - rotated = rotate_fn({"img": im_type(self.imt[0]), "seg": im_type(self.segn[0])}) + im = im_type(self.imt[0]) + rotated = rotate_fn({"img": im, "seg": im_type(self.segn[0])}) if keep_size: np.testing.assert_allclose(self.imt[0].shape, rotated["img"].shape) _order = 0 if mode == "nearest" else 1 @@ -60,11 +62,14 @@ def test_correct_results(self, im_type, angle, keep_size, mode, padding_mode, al rotated[k] = v.cpu() if isinstance(v, torch.Tensor) else v good = np.sum(np.isclose(expected, rotated["img"][0], atol=1e-3)) self.assertLessEqual(np.abs(good - expected.size), 5, "diff at most 5 pixels") + test_local_inversion(rotate_fn, rotated, {"img": im}, "img") expected = scipy.ndimage.rotate( self.segn[0, 0], -np.rad2deg(angle), (0, 1), not keep_size, order=0, mode=_mode, prefilter=False ) expected = np.stack(expected).astype(int) + if isinstance(rotated["seg"], MetaTensor): + rotated["seg"] = rotated["seg"].as_tensor() # pytorch 1.7 compatible self.assertLessEqual(np.count_nonzero(expected != rotated["seg"][0]), 30) @@ -96,6 +101,8 @@ def test_correct_results(self, im_type, angle, keep_size, mode, padding_mode, al self.segn[0, 0], np.rad2deg(angle), (0, 2), not keep_size, order=0, mode=_mode, prefilter=False ) expected = np.stack(expected).astype(int) + if isinstance(rotated["seg"], MetaTensor): + rotated["seg"] = rotated["seg"].as_tensor() # pytorch 1.7 compatible self.assertLessEqual(np.count_nonzero(expected != rotated["seg"][0]), 160) @@ -127,6 +134,8 @@ def test_correct_results(self, im_type, angle, keep_size, mode, padding_mode, al self.segn[0, 0], -np.rad2deg(angle), (0, 1), not keep_size, order=0, mode=_mode, prefilter=False ) expected = np.stack(expected).astype(int) + if isinstance(rotated["seg"], MetaTensor): + rotated["seg"] = rotated["seg"].as_tensor() # pytorch 1.7 compatible self.assertLessEqual(np.count_nonzero(expected != rotated["seg"][0]), 160) diff --git a/tests/test_save_image.py b/tests/test_save_image.py index a1297c1e61..6591283c22 100644 --- a/tests/test_save_image.py +++ b/tests/test_save_image.py @@ -13,10 +13,10 @@ import tempfile import unittest -import numpy as np import torch from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms import SaveImage TEST_CASE_1 = [torch.randint(0, 255, (1, 2, 3, 4)), {"filename_or_obj": "testfile0.nii.gz"}, ".nii.gz", False] @@ -26,7 +26,7 @@ TEST_CASE_3 = [torch.randint(0, 255, (1, 2, 3, 4)), {"filename_or_obj": "testfile0.nrrd"}, ".nrrd", False] TEST_CASE_4 = [ - np.random.randint(0, 255, (3, 2, 4, 5), dtype=np.uint8), + torch.randint(0, 255, (3, 2, 4, 5), dtype=torch.uint8), {"filename_or_obj": "testfile0.dcm"}, ".dcm", False, @@ -36,6 +36,9 @@ class TestSaveImage(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) def test_saved_content(self, test_data, meta_data, output_ext, resample): + if meta_data is not None: + test_data = MetaTensor(test_data, meta=meta_data) + with tempfile.TemporaryDirectory() as tempdir: trans = SaveImage( output_dir=tempdir, @@ -43,7 +46,7 @@ def test_saved_content(self, test_data, meta_data, output_ext, resample): resample=resample, separate_folder=False, # test saving into the same folder ) - trans(test_data, meta_data) + trans(test_data) filepath = "testfile0" if meta_data is not None else "0" self.assertTrue(os.path.exists(os.path.join(tempdir, filepath + "_trans" + output_ext))) diff --git a/tests/test_save_imaged.py b/tests/test_save_imaged.py index a6988683e5..96b6fb1626 100644 --- a/tests/test_save_imaged.py +++ b/tests/test_save_imaged.py @@ -16,19 +16,18 @@ import torch from parameterized import parameterized +from monai.data.meta_tensor import MetaTensor from monai.transforms import SaveImaged -from monai.utils.enums import PostFix TEST_CASE_1 = [ - {"img": torch.randint(0, 255, (1, 2, 3, 4)), PostFix.meta("img"): {"filename_or_obj": "testfile0.nii.gz"}}, + {"img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nii.gz"})}, ".nii.gz", False, ] TEST_CASE_2 = [ { - "img": torch.randint(0, 255, (1, 2, 3, 4)), - PostFix.meta("img"): {"filename_or_obj": "testfile0.nii.gz"}, + "img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nii.gz"}), "patch_index": 6, }, ".nii.gz", @@ -37,8 +36,7 @@ TEST_CASE_3 = [ { - "img": torch.randint(0, 255, (1, 2, 3, 4)), - PostFix.meta("img"): {"filename_or_obj": "testfile0.nrrd"}, + "img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nrrd"}), "patch_index": 6, }, ".nrrd", @@ -52,7 +50,6 @@ def test_saved_content(self, test_data, output_ext, resample): with tempfile.TemporaryDirectory() as tempdir: trans = SaveImaged( keys=["img", "pred"], - meta_keys=PostFix.meta("img"), output_dir=tempdir, output_ext=output_ext, resample=resample, @@ -60,7 +57,7 @@ def test_saved_content(self, test_data, output_ext, resample): ) trans(test_data) - patch_index = test_data[PostFix.meta("img")].get("patch_index", None) + patch_index = test_data["img"].meta.get("patch_index", None) patch_index = f"_{patch_index}" if patch_index is not None else "" filepath = os.path.join("testfile0", "testfile0" + "_trans" + patch_index + output_ext) self.assertTrue(os.path.exists(os.path.join(tempdir, filepath))) diff --git a/tests/test_savitzky_golay_smooth.py b/tests/test_savitzky_golay_smooth.py index ac42cf806e..b296372986 100644 --- a/tests/test_savitzky_golay_smooth.py +++ b/tests/test_savitzky_golay_smooth.py @@ -12,11 +12,10 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms import SavitzkyGolaySmooth -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_NDARRAYS, assert_allclose # Zero-padding trivial tests @@ -65,7 +64,7 @@ class TestSavitzkyGolaySmooth(unittest.TestCase): def test_value(self, arguments, image, expected_data, atol): for p in TEST_NDARRAYS: result = SavitzkyGolaySmooth(**arguments)(p(image.astype(np.float32))) - torch.testing.assert_allclose(result, p(expected_data.astype(np.float32)), rtol=1e-4, atol=atol) + assert_allclose(result, p(expected_data.astype(np.float32)), rtol=1e-4, atol=atol, type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_scale_intensity.py b/tests/test_scale_intensity.py index bd1adac4f4..9941172b0e 100644 --- a/tests/test_scale_intensity.py +++ b/tests/test_scale_intensity.py @@ -12,6 +12,7 @@ import unittest import numpy as np +from parameterized import parameterized from monai.transforms import ScaleIntensity from monai.transforms.utils import rescale_array @@ -19,29 +20,30 @@ class TestScaleIntensity(NumpyImageTestCase2D): - def test_range_scale(self): - for p in TEST_NDARRAYS: - scaler = ScaleIntensity(minv=1.0, maxv=2.0) - result = scaler(p(self.imt)) - mina = self.imt.min() - maxa = self.imt.max() - norm = (self.imt - mina) / (maxa - mina) - expected = p((norm * (2.0 - 1.0)) + 1.0) - assert_allclose(result, expected, type_test=False, rtol=1e-7, atol=0) + @parameterized.expand([[p] for p in TEST_NDARRAYS]) + def test_range_scale(self, p): + scaler = ScaleIntensity(minv=1.0, maxv=2.0) + im = p(self.imt) + result = scaler(im) + mina = self.imt.min() + maxa = self.imt.max() + norm = (self.imt - mina) / (maxa - mina) + expected = p((norm * (2.0 - 1.0)) + 1.0) + assert_allclose(result, expected, type_test="tensor", rtol=1e-7, atol=0) def test_factor_scale(self): for p in TEST_NDARRAYS: scaler = ScaleIntensity(minv=None, maxv=None, factor=0.1) result = scaler(p(self.imt)) expected = p((self.imt * (1 + 0.1)).astype(np.float32)) - assert_allclose(result, p(expected), rtol=1e-7, atol=0) + assert_allclose(result, p(expected), type_test="tensor", rtol=1e-7, atol=0) def test_max_none(self): for p in TEST_NDARRAYS: scaler = ScaleIntensity(minv=0.0, maxv=None, factor=0.1) result = scaler(p(self.imt)) expected = rescale_array(p(self.imt), minv=0.0, maxv=None) - assert_allclose(result, expected, rtol=1e-3, atol=1e-3) + assert_allclose(result, expected, type_test="tensor", rtol=1e-3, atol=1e-3) def test_int(self): """integers should be handled by converting them to floats first.""" @@ -53,7 +55,7 @@ def test_int(self): maxa = _imt.max() norm = (_imt - mina) / (maxa - mina) expected = p((norm * (2.0 - 1.0)) + 1.0) - assert_allclose(result, expected, type_test=False, rtol=1e-7, atol=0) + assert_allclose(result, expected, type_test="tensor", rtol=1e-7, atol=0) def test_channel_wise(self): for p in TEST_NDARRAYS: @@ -65,7 +67,7 @@ def test_channel_wise(self): for i, c in enumerate(data): norm = (c - mina) / (maxa - mina) expected = p((norm * (2.0 - 1.0)) + 1.0) - assert_allclose(result[i], expected, type_test=False, rtol=1e-7, atol=0) + assert_allclose(result[i], expected, type_test="tensor", rtol=1e-7, atol=0) if __name__ == "__main__": diff --git a/tests/test_scale_intensity_range.py b/tests/test_scale_intensity_range.py index faddf9001b..958881f790 100644 --- a/tests/test_scale_intensity_range.py +++ b/tests/test_scale_intensity_range.py @@ -24,7 +24,7 @@ def test_image_scale_intensity_range(self): scaled = scaler(p(self.imt)) self.assertTrue(scaled.dtype, np.uint8) expected = (((self.imt - 20) / 88) * 30 + 50).astype(np.uint8) - assert_allclose(scaled, p(expected)) + assert_allclose(scaled, p(expected), type_test="tensor") def test_image_scale_intensity_range_none_clip(self): scaler = ScaleIntensityRange(a_min=20, a_max=108, b_min=None, b_max=80, clip=True, dtype=np.uint8) @@ -32,7 +32,7 @@ def test_image_scale_intensity_range_none_clip(self): scaled = scaler(p(self.imt)) self.assertTrue(scaled.dtype, np.uint8) expected = (np.clip((self.imt - 20) / 88, None, 80)).astype(np.uint8) - assert_allclose(scaled, p(expected)) + assert_allclose(scaled, p(expected), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_scale_intensity_range_percentiles.py b/tests/test_scale_intensity_range_percentiles.py index f8656dd929..184e1dff0c 100644 --- a/tests/test_scale_intensity_range_percentiles.py +++ b/tests/test_scale_intensity_range_percentiles.py @@ -31,7 +31,7 @@ def test_scaling(self): scaler = ScaleIntensityRangePercentiles(lower=lower, upper=upper, b_min=b_min, b_max=b_max, dtype=np.uint8) for p in TEST_NDARRAYS: result = scaler(p(img)) - assert_allclose(result, p(expected), rtol=1e-4) + assert_allclose(result, p(expected), type_test="tensor", rtol=1e-4) def test_relative_scaling(self): img = self.imt[0] @@ -50,14 +50,16 @@ def test_relative_scaling(self): for p in TEST_NDARRAYS: result = scaler(p(img)) - assert_allclose(result, p(expected_img), rtol=1e-3) + assert_allclose(result, p(expected_img), type_test="tensor", rtol=1e-3) scaler = ScaleIntensityRangePercentiles( lower=lower, upper=upper, b_min=b_min, b_max=b_max, relative=True, clip=True ) for p in TEST_NDARRAYS: result = scaler(p(img)) - assert_allclose(result, p(np.clip(expected_img, expected_b_min, expected_b_max)), rtol=1e-4) + assert_allclose( + result, p(np.clip(expected_img, expected_b_min, expected_b_max)), type_test="tensor", rtol=1e-4 + ) def test_invalid_instantiation(self): self.assertRaises(ValueError, ScaleIntensityRangePercentiles, lower=-10, upper=99, b_min=0, b_max=255) @@ -83,7 +85,7 @@ def test_channel_wise(self): for p in TEST_NDARRAYS: result = scaler(p(img)) - assert_allclose(result, p(expected), rtol=1e-4) + assert_allclose(result, p(expected), type_test="tensor", rtol=1e-4) if __name__ == "__main__": diff --git a/tests/test_scale_intensity_range_percentilesd.py b/tests/test_scale_intensity_range_percentilesd.py index 5441832a77..b35f937b95 100644 --- a/tests/test_scale_intensity_range_percentilesd.py +++ b/tests/test_scale_intensity_range_percentilesd.py @@ -34,12 +34,11 @@ def test_scaling(self): scaler = ScaleIntensityRangePercentilesd( keys=data.keys(), lower=lower, upper=upper, b_min=b_min, b_max=b_max, dtype=np.uint8 ) - assert_allclose(p(expected), scaler(data)["img"], rtol=1e-4) + assert_allclose(scaler(data)["img"], p(expected), type_test="tensor", rtol=1e-4) def test_relative_scaling(self): img = self.imt - data = {} - data["img"] = img + data = {"img": img} lower = 10 upper = 99 b_min = 100 @@ -55,7 +54,7 @@ def test_relative_scaling(self): expected_img = (img - expected_a_min) / (expected_a_max - expected_a_min) expected_img = (expected_img * (expected_b_max - expected_b_min)) + expected_b_min - np.testing.assert_allclose(expected_img, scaler(data)["img"]) + np.testing.assert_allclose(expected_img, scaler(data)["img"], rtol=1e-3, atol=0.1) def test_invalid_instantiation(self): self.assertRaises( @@ -92,7 +91,7 @@ def test_channel_wise(self): for p in TEST_NDARRAYS: data = {"img": p(img)} - assert_allclose(scaler(data)["img"], p(expected), rtol=1e-4) + assert_allclose(scaler(data)["img"], p(expected), type_test="tensor", rtol=1e-4) if __name__ == "__main__": diff --git a/tests/test_scale_intensity_ranged.py b/tests/test_scale_intensity_ranged.py index ffbd3e44c4..4ac4910e37 100644 --- a/tests/test_scale_intensity_ranged.py +++ b/tests/test_scale_intensity_ranged.py @@ -23,7 +23,7 @@ def test_image_scale_intensity_ranged(self): scaled = scaler({key: p(self.imt)}) expected = (self.imt - 20) / 88 expected = expected * 30 + 50 - assert_allclose(scaled[key], p(expected)) + assert_allclose(scaled[key], p(expected), type_test="tensor") def test_image_scale_intensity_ranged_none(self): key = "img" @@ -31,7 +31,7 @@ def test_image_scale_intensity_ranged_none(self): for p in TEST_NDARRAYS: scaled = scaler({key: p(self.imt)}) expected = (self.imt - 20) / 88 - assert_allclose(scaled[key], p(expected)) + assert_allclose(scaled[key], p(expected), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_scale_intensityd.py b/tests/test_scale_intensityd.py index 42f1527490..d560523214 100644 --- a/tests/test_scale_intensityd.py +++ b/tests/test_scale_intensityd.py @@ -27,7 +27,7 @@ def test_range_scale(self): maxa = np.max(self.imt) norm = (self.imt - mina) / (maxa - mina) expected = (norm * (2.0 - 1.0)) + 1.0 - assert_allclose(result[key], p(expected)) + assert_allclose(result[key], p(expected), type_test="tensor") def test_factor_scale(self): key = "img" @@ -35,7 +35,7 @@ def test_factor_scale(self): scaler = ScaleIntensityd(keys=[key], minv=None, maxv=None, factor=0.1) result = scaler({key: p(self.imt)}) expected = (self.imt * (1 + 0.1)).astype(np.float32) - assert_allclose(result[key], p(expected)) + assert_allclose(result[key], p(expected), type_test="tensor") def test_channel_wise(self): key = "img" @@ -48,7 +48,7 @@ def test_channel_wise(self): for i, c in enumerate(data): norm = (c - mina) / (maxa - mina) expected = p((norm * (2.0 - 1.0)) + 1.0) - assert_allclose(result[key][i], expected, type_test=False, rtol=1e-7, atol=0) + assert_allclose(result[key][i], expected, type_test="tensor", rtol=1e-7, atol=0) if __name__ == "__main__": diff --git a/tests/test_senet.py b/tests/test_senet.py index 57ca49237d..80d2b071c3 100644 --- a/tests/test_senet.py +++ b/tests/test_senet.py @@ -65,7 +65,15 @@ def test_script(self, net, net_args): class TestPretrainedSENET(unittest.TestCase): def setUp(self): self.original_urls = se_mod.SE_NET_MODELS.copy() - if test_is_quick(): + replace_url = test_is_quick() + if not replace_url: + try: + SEResNet50(pretrained=True, spatial_dims=2, in_channels=3, num_classes=2) + except OSError as rt_e: + print(rt_e) + if "certificate" in str(rt_e): # [SSL: CERTIFICATE_VERIFY_FAILED] + replace_url = True + if replace_url: testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") testing_data_urls = { "senet154": { diff --git a/tests/test_shift_intensityd.py b/tests/test_shift_intensityd.py index e28b7f54e4..b5a2a3218d 100644 --- a/tests/test_shift_intensityd.py +++ b/tests/test_shift_intensityd.py @@ -25,7 +25,7 @@ def test_value(self): shifter = ShiftIntensityd(keys=[key], offset=1.0) result = shifter({key: p(self.imt)}) expected = self.imt + 1.0 - assert_allclose(result[key], p(expected)) + assert_allclose(result[key], p(expected), type_test="tensor") def test_factor(self): key = "img" diff --git a/tests/test_shuffle_buffer.py b/tests/test_shuffle_buffer.py new file mode 100644 index 0000000000..40012fbf93 --- /dev/null +++ b/tests/test_shuffle_buffer.py @@ -0,0 +1,34 @@ +# 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. + +import sys +import unittest + +import numpy as np + +from monai.data import DataLoader, ShuffleBuffer +from monai.utils import convert_data_type + + +class TestShuffleBuffer(unittest.TestCase): + def test_shape(self): + buffer = ShuffleBuffer([1, 2, 3, 4], seed=0) + num_workers = 2 if sys.platform == "linux" else 0 + dataloader = DataLoader(dataset=buffer, batch_size=2, num_workers=num_workers) + output = [convert_data_type(x, np.ndarray)[0] for x in dataloader] + if num_workers == 0: + np.testing.assert_allclose(output, [[2, 1], [3, 4]]) + else: # multiprocess shuffle + np.testing.assert_allclose(output, [[2, 3], [1, 4]]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_slice_inferer.py b/tests/test_slice_inferer.py index 3c52082d85..0f33385f42 100644 --- a/tests/test_slice_inferer.py +++ b/tests/test_slice_inferer.py @@ -46,6 +46,9 @@ def test_shape(self, spatial_dim): self.assertTupleEqual(result.shape, input_volume.shape) + # test that the inferer can be run multiple times + result = inferer(input_volume, model) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_sliding_patch_wsi_dataset.py b/tests/test_sliding_patch_wsi_dataset.py new file mode 100644 index 0000000000..d639d000c5 --- /dev/null +++ b/tests/test_sliding_patch_wsi_dataset.py @@ -0,0 +1,269 @@ +# 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. + +import os +import unittest +from unittest import skipUnless + +import numpy as np +from numpy.testing import assert_array_equal +from parameterized import parameterized + +from monai.data import SlidingPatchWSIDataset +from monai.utils import WSIPatchKeys, optional_import, set_determinism +from tests.utils import download_url_or_skip_test, testing_data_config + +set_determinism(0) + +cucim, has_cucim = optional_import("cucim") +has_cucim = has_cucim and hasattr(cucim, "CuImage") +openslide, has_osl = optional_import("openslide") +imwrite, has_tiff = optional_import("tifffile", name="imwrite") +_, has_codec = optional_import("imagecodecs") +has_tiff = has_tiff and has_codec + + +FILE_KEY = "wsi_img" +FILE_URL = testing_data_config("images", FILE_KEY, "url") +base_name, extension = os.path.basename(f"{FILE_URL}"), ".tiff" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", "temp_" + base_name + extension) + +FILE_PATH_SMALL_0 = os.path.join(os.path.dirname(__file__), "testing_data", "temp_wsi_inference_0.tiff") +FILE_PATH_SMALL_1 = os.path.join(os.path.dirname(__file__), "testing_data", "temp_wsi_inference_1.tiff") +ARRAY_SMALL_0 = np.random.randint(low=0, high=255, size=(3, 4, 4), dtype=np.uint8) +ARRAY_SMALL_1 = np.random.randint(low=0, high=255, size=(3, 5, 5), dtype=np.uint8) + +TEST_CASE_SMALL_0 = [ + {"data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0}], "patch_size": (2, 2)}, + [ + {"image": ARRAY_SMALL_0[:, :2, :2]}, + {"image": ARRAY_SMALL_0[:, :2, 2:]}, + {"image": ARRAY_SMALL_0[:, 2:, :2]}, + {"image": ARRAY_SMALL_0[:, 2:, 2:]}, + ], +] + +TEST_CASE_SMALL_1 = [ + {"data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (2, 2)}]}, + [ + {"image": ARRAY_SMALL_0[:, :2, :2]}, + {"image": ARRAY_SMALL_0[:, :2, 2:]}, + {"image": ARRAY_SMALL_0[:, 2:, :2]}, + {"image": ARRAY_SMALL_0[:, 2:, 2:]}, + ], +] + +TEST_CASE_SMALL_2 = [ + {"data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0}], "patch_size": (2, 2), "overlap": 0.5}, + [ + {"image": ARRAY_SMALL_0[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_0[:, 0:2, 1:3]}, + {"image": ARRAY_SMALL_0[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_0[:, 1:3, 0:2]}, + {"image": ARRAY_SMALL_0[:, 1:3, 1:3]}, + {"image": ARRAY_SMALL_0[:, 1:3, 2:4]}, + {"image": ARRAY_SMALL_0[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_0[:, 2:4, 1:3]}, + {"image": ARRAY_SMALL_0[:, 2:4, 2:4]}, + ], +] + +TEST_CASE_SMALL_3 = [ + {"data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0}], "patch_size": (3, 3), "overlap": 2.0 / 3.0}, + [ + {"image": ARRAY_SMALL_0[:, :3, :3]}, + {"image": ARRAY_SMALL_0[:, :3, 1:]}, + {"image": ARRAY_SMALL_0[:, 1:, :3]}, + {"image": ARRAY_SMALL_0[:, 1:, 1:]}, + ], +] + +TEST_CASE_SMALL_4 = [ + { + "data": [ + {"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0}, + {"image": FILE_PATH_SMALL_1, WSIPatchKeys.LEVEL: 0}, + ], + "patch_size": (2, 2), + }, + [ + {"image": ARRAY_SMALL_0[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_0[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_0[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_0[:, 2:4, 2:4]}, + {"image": ARRAY_SMALL_1[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_1[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_1[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_1[:, 2:4, 2:4]}, + ], +] + +TEST_CASE_SMALL_5 = [ + { + "data": [ + {"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (2, 2)}, + {"image": FILE_PATH_SMALL_1, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (3, 3)}, + ] + }, + [ + {"image": ARRAY_SMALL_0[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_0[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_0[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_0[:, 2:4, 2:4]}, + {"image": ARRAY_SMALL_1[:, 0:3, 0:3]}, + ], +] + +TEST_CASE_SMALL_6 = [ + { + "data": [ + {"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 1, WSIPatchKeys.SIZE: (1, 1)}, + {"image": FILE_PATH_SMALL_1, WSIPatchKeys.LEVEL: 2, WSIPatchKeys.SIZE: (4, 4)}, + ], + "patch_size": (2, 2), + "patch_level": 0, + }, + [ + {"image": ARRAY_SMALL_0[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_0[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_0[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_0[:, 2:4, 2:4]}, + {"image": ARRAY_SMALL_1[:, 0:2, 0:2]}, + {"image": ARRAY_SMALL_1[:, 0:2, 2:4]}, + {"image": ARRAY_SMALL_1[:, 2:4, 0:2]}, + {"image": ARRAY_SMALL_1[:, 2:4, 2:4]}, + ], +] + + +TEST_CASE_SMALL_7 = [ + {"data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (2, 2)}], "offset": (1, 0)}, + [{"image": ARRAY_SMALL_0[:, 1:3, :2]}, {"image": ARRAY_SMALL_0[:, 1:3, 2:]}], +] + +TEST_CASE_SMALL_8 = [ + { + "data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (2, 2)}], + "offset": "random", + "offset_limits": (0, 2), + }, + [{"image": ARRAY_SMALL_0[:, 1:3, :2]}, {"image": ARRAY_SMALL_0[:, 1:3, 2:]}], +] + +TEST_CASE_SMALL_9 = [ + { + "data": [{"image": FILE_PATH_SMALL_0, WSIPatchKeys.LEVEL: 0, WSIPatchKeys.SIZE: (2, 2)}], + "offset": "random", + "offset_limits": ((0, 3), (0, 2)), + }, + [{"image": ARRAY_SMALL_0[:, :2, 1:3]}, {"image": ARRAY_SMALL_0[:, 2:, 1:3]}], +] + +TEST_CASE_LARGE_0 = [ + {"data": [{"image": FILE_PATH, WSIPatchKeys.LEVEL: 8, WSIPatchKeys.SIZE: (64, 50)}]}, + [ + {"step_loc": (0, 0), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (0, 1), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (0, 2), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 0), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 1), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 2), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + ], +] + +TEST_CASE_LARGE_1 = [ + { + "data": [ + {"image": FILE_PATH, WSIPatchKeys.LEVEL: 8, WSIPatchKeys.SIZE: (64, 50)}, + {"image": FILE_PATH, WSIPatchKeys.LEVEL: 7, WSIPatchKeys.SIZE: (125, 110)}, + ] + }, + [ + {"step_loc": (0, 0), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (0, 1), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (0, 2), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 0), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 1), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (1, 2), "patch_size": (64, 50), "patch_level": 8, "ratio": 257.06195068359375}, + {"step_loc": (0, 0), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + {"step_loc": (0, 1), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + {"step_loc": (0, 2), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + {"step_loc": (1, 0), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + {"step_loc": (1, 1), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + {"step_loc": (1, 2), "patch_size": (125, 110), "patch_level": 7, "ratio": 128.10186767578125}, + ], +] + + +@skipUnless(has_cucim or has_tiff, "Requires cucim, openslide, or tifffile!") +def setUpModule(): + for info in [(ARRAY_SMALL_0, FILE_PATH_SMALL_0), (ARRAY_SMALL_1, FILE_PATH_SMALL_1)]: + array = info[0].transpose([1, 2, 0]) + imwrite(info[1], array, shape=array.shape, photometric="rgb") + hash_type = testing_data_config("images", FILE_KEY, "hash_type") + hash_val = testing_data_config("images", FILE_KEY, "hash_val") + download_url_or_skip_test(FILE_URL, FILE_PATH, hash_type=hash_type, hash_val=hash_val) + + +class SlidingPatchWSIDatasetTests: + class Tests(unittest.TestCase): + backend = None + + @parameterized.expand( + [ + TEST_CASE_SMALL_0, + TEST_CASE_SMALL_1, + TEST_CASE_SMALL_2, + TEST_CASE_SMALL_3, + TEST_CASE_SMALL_4, + TEST_CASE_SMALL_5, + TEST_CASE_SMALL_6, + TEST_CASE_SMALL_7, + TEST_CASE_SMALL_8, + TEST_CASE_SMALL_9, + ] + ) + def test_read_patches(self, input_parameters, expected): + if self.backend == "openslide": + return + dataset = SlidingPatchWSIDataset(reader=self.backend, **input_parameters) + self.assertEqual(len(dataset), len(expected)) + for i, sample in enumerate(dataset): + self.assertTupleEqual(sample["image"].shape, expected[i]["image"].shape) + + @parameterized.expand([TEST_CASE_LARGE_0, TEST_CASE_LARGE_1]) + def test_read_patches_large(self, input_parameters, expected): + dataset = SlidingPatchWSIDataset(reader=self.backend, **input_parameters) + self.assertEqual(len(dataset), len(expected)) + for i, sample in enumerate(dataset): + self.assertEqual(sample["metadata"][WSIPatchKeys.LEVEL], expected[i]["patch_level"]) + assert_array_equal(sample["metadata"][WSIPatchKeys.SIZE], expected[i]["patch_size"]) + steps = [round(expected[i]["ratio"] * s) for s in expected[i]["patch_size"]] + expected_location = tuple(expected[i]["step_loc"][j] * steps[j] for j in range(len(steps))) + assert_array_equal(sample["metadata"][WSIPatchKeys.LOCATION], expected_location) + + +@skipUnless(has_cucim, "Requires cucim") +class TestSlidingPatchWSIDatasetCuCIM(SlidingPatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "cucim" + + +@skipUnless(has_osl, "Requires openslide") +class TestSlidingPatchWSIDatasetOpenSlide(SlidingPatchWSIDatasetTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "openslide" + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_sliding_window_inference.py b/tests/test_sliding_window_inference.py index 5b6995c1ea..8b8ec47d32 100644 --- a/tests/test_sliding_window_inference.py +++ b/tests/test_sliding_window_inference.py @@ -15,9 +15,10 @@ import torch from parameterized import parameterized +from monai.data.utils import list_data_collate from monai.inferers import SlidingWindowInferer, sliding_window_inference from monai.utils import optional_import -from tests.utils import skip_if_no_cuda +from tests.utils import TEST_TORCH_AND_META_TENSORS, skip_if_no_cuda _, has_tqdm = optional_import("tqdm") @@ -68,9 +69,11 @@ def compute(data): np.testing.assert_string_equal(device.type, result.device.type) np.testing.assert_allclose(result.cpu().numpy(), expected_val) - def test_default_device(self): + @parameterized.expand([[x] for x in TEST_TORCH_AND_META_TENSORS]) + def test_default_device(self, data_type): device = "cuda" if torch.cuda.is_available() else "cpu:0" - inputs = torch.ones((1, 3, 16, 15, 7)).to(device=device) + inputs = data_type(torch.ones((3, 16, 15, 7))).to(device=device) + inputs = list_data_collate([inputs]) # make a proper batch roi_shape = (4, 10, 7) sw_batch_size = 10 @@ -82,9 +85,11 @@ def compute(data): expected_val = np.ones((1, 3, 16, 15, 7), dtype=np.float32) + 1 np.testing.assert_allclose(result.cpu().numpy(), expected_val) + @parameterized.expand([[x] for x in TEST_TORCH_AND_META_TENSORS]) @skip_if_no_cuda - def test_sw_device(self): - inputs = torch.ones((1, 3, 16, 15, 7)).to(device="cpu") + def test_sw_device(self, data_type): + inputs = data_type(torch.ones((3, 16, 15, 7))).to(device="cpu") + inputs = list_data_collate([inputs]) # make a proper batch roi_shape = (4, 10, 7) sw_batch_size = 10 @@ -177,6 +182,11 @@ def compute(self, data): ) np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4) + result = SlidingWindowInferer( + roi_shape, sw_batch_size, overlap=0.5, mode="gaussian", sigma_scale=[1.0, 1.0], cache_roi_weight_map=True + )(inputs, _Pred().compute) + np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4) + def test_cval(self): device = "cuda" if torch.cuda.is_available() else "cpu:0" inputs = torch.ones((1, 1, 3, 3)).to(device=device) @@ -227,6 +237,7 @@ def compute(data, test1, test2): device, device, has_tqdm, + None, t1, test2=t2, ) @@ -238,6 +249,67 @@ def compute(data, test1, test2): )(inputs, compute, t1, test2=t2) np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4) + def test_multioutput(self): + device = "cuda" if torch.cuda.is_available() else "cpu:0" + inputs = torch.ones((1, 6, 20, 20)).to(device=device) + roi_shape = (8, 8) + sw_batch_size = 10 + + def compute(data): + return data + 1, data[:, ::3, ::2, ::2] + 2, data[:, ::2, ::4, ::4] + 3 + + def compute_dict(data): + return {1: data + 1, 2: data[:, ::3, ::2, ::2] + 2, 3: data[:, ::2, ::4, ::4] + 3} + + result = sliding_window_inference( + inputs, + roi_shape, + sw_batch_size, + compute, + 0.5, + "constant", + 1.0, + "constant", + 0.0, + device, + device, + has_tqdm, + None, + ) + result_dict = sliding_window_inference( + inputs, + roi_shape, + sw_batch_size, + compute_dict, + 0.5, + "constant", + 1.0, + "constant", + 0.0, + device, + device, + has_tqdm, + None, + ) + expected = (np.ones((1, 6, 20, 20)) + 1, np.ones((1, 2, 10, 10)) + 2, np.ones((1, 3, 5, 5)) + 3) + expected_dict = {1: np.ones((1, 6, 20, 20)) + 1, 2: np.ones((1, 2, 10, 10)) + 2, 3: np.ones((1, 3, 5, 5)) + 3} + for rr, ee in zip(result, expected): + np.testing.assert_allclose(rr.cpu().numpy(), ee, rtol=1e-4) + for rr, _ in zip(result_dict, expected_dict): + np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4) + + result = SlidingWindowInferer( + roi_shape, sw_batch_size, overlap=0.5, mode="constant", cval=-1, progress=has_tqdm + )(inputs, compute) + for rr, ee in zip(result, expected): + np.testing.assert_allclose(rr.cpu().numpy(), ee, rtol=1e-4) + + result_dict = SlidingWindowInferer( + roi_shape, sw_batch_size, overlap=0.5, mode="constant", cval=-1, progress=has_tqdm + )(inputs, compute_dict) + for rr, _ in zip(result_dict, expected_dict): + np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_smartcachedataset.py b/tests/test_smartcachedataset.py index e7d51be63a..9f9043d19e 100644 --- a/tests/test_smartcachedataset.py +++ b/tests/test_smartcachedataset.py @@ -17,10 +17,12 @@ import nibabel as nib import numpy as np +import torch from parameterized import parameterized from monai.data import DataLoader, SmartCacheDataset from monai.transforms import Compose, Lambda, LoadImaged +from tests.utils import assert_allclose TEST_CASE_1 = [0.1, 0, Compose([LoadImaged(keys=["image", "label", "extra"])])] @@ -56,6 +58,17 @@ def test_shape(self, replace_rate, num_replace_workers, transform): num_init_workers=4, num_replace_workers=num_replace_workers, ) + if transform is None: + # Check without providing transfrom + dataset2 = SmartCacheDataset( + data=test_data, + replace_rate=replace_rate, + cache_num=16, + num_init_workers=4, + num_replace_workers=num_replace_workers, + ) + for k in ["image", "label", "extra"]: + self.assertEqual(dataset[0][k], dataset2[0][k]) self.assertEqual(len(dataset._cache), dataset.cache_num) for i in range(dataset.cache_num): @@ -66,8 +79,8 @@ def test_shape(self, replace_rate, num_replace_workers, transform): for _ in range(3): dataset.update_cache() self.assertIsNotNone(dataset[15]) - if isinstance(dataset[15]["image"], np.ndarray): - np.testing.assert_allclose(dataset[15]["image"], dataset[15]["label"]) + if isinstance(dataset[15]["image"], (np.ndarray, torch.Tensor)): + assert_allclose(dataset[15]["image"], dataset[15]["label"]) else: self.assertIsInstance(dataset[15]["image"], str) dataset.shutdown() diff --git a/tests/test_smooth_field.py b/tests/test_smooth_field.py index 5849b96167..c67865ba39 100644 --- a/tests/test_smooth_field.py +++ b/tests/test_smooth_field.py @@ -87,7 +87,7 @@ def test_rand_smooth_field_adjust_contrastd(self, input_param, input_data, expec res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") def test_rand_smooth_field_adjust_contrastd_pad(self): input_param, input_data, expected_val = TESTS_CONTRAST[0] @@ -98,7 +98,7 @@ def test_rand_smooth_field_adjust_contrastd_pad(self): res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") @parameterized.expand(TESTS_INTENSITY) def test_rand_smooth_field_adjust_intensityd(self, input_param, input_data, expected_val): @@ -108,7 +108,7 @@ def test_rand_smooth_field_adjust_intensityd(self, input_param, input_data, expe res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") def test_rand_smooth_field_adjust_intensityd_pad(self): input_param, input_data, expected_val = TESTS_INTENSITY[0] @@ -119,7 +119,7 @@ def test_rand_smooth_field_adjust_intensityd_pad(self): res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") @parameterized.expand(TESTS_DEFORM) def test_rand_smooth_deformd(self, input_param, input_data, expected_val): @@ -129,7 +129,7 @@ def test_rand_smooth_deformd(self, input_param, input_data, expected_val): res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") def test_rand_smooth_deformd_pad(self): input_param, input_data, expected_val = TESTS_DEFORM[0] @@ -140,4 +140,8 @@ def test_rand_smooth_deformd_pad(self): res = g(input_data) for key, result in res.items(): expected = expected_val[key] - assert_allclose(result, expected, rtol=_rtol, atol=1e-1) + assert_allclose(result, expected, rtol=_rtol, atol=1e-1, type_test="tensor") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_spacing.py b/tests/test_spacing.py index 80df981b73..6e56a21b63 100644 --- a/tests/test_spacing.py +++ b/tests/test_spacing.py @@ -15,127 +15,140 @@ import torch from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor from monai.data.utils import affine_to_spacing from monai.transforms import Spacing from monai.utils import ensure_tuple, fall_back_tuple -from tests.utils import TEST_NDARRAYS +from tests.utils import TEST_DEVICES, assert_allclose TESTS = [] -for p in TEST_NDARRAYS: +for device in TEST_DEVICES: TESTS.append( [ - p, {"pixdim": (1.0, 1.5), "padding_mode": "zeros", "dtype": float}, - np.arange(4).reshape((1, 2, 2)) + 1.0, # data - {"affine": np.eye(4)}, - np.array([[[1.0, 1.0], [3.0, 2.0]]]), + torch.arange(4).reshape((1, 2, 2)) + 1.0, # data + torch.eye(4), + {}, + torch.tensor([[[1.0, 1.0], [3.0, 2.0]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": 1.0, "padding_mode": "zeros", "dtype": float}, - np.ones((1, 2, 1, 2)), # data - {"affine": np.eye(4)}, - np.array([[[[1.0, 1.0]], [[1.0, 1.0]]]]), + torch.ones((1, 2, 1, 2)), # data + torch.eye(4), + {}, + torch.tensor([[[[1.0, 1.0]], [[1.0, 1.0]]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": (1.0, 1.0, 1.0), "padding_mode": "zeros", "dtype": float}, - np.ones((1, 2, 1, 2)), # data - {"affine": np.eye(4)}, - np.array([[[[1.0, 1.0]], [[1.0, 1.0]]]]), + torch.ones((1, 2, 1, 2)), # data + torch.eye(4), + {}, + torch.tensor([[[[1.0, 1.0]], [[1.0, 1.0]]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": (1.0, 0.2, 1.5), "diagonal": False, "padding_mode": "zeros", "align_corners": True}, - np.ones((1, 2, 1, 2)), # data - {"affine": np.array([[2, 1, 0, 4], [-1, -3, 0, 5], [0, 0, 2.0, 5], [0, 0, 0, 1]])}, - np.array([[[[0.95527864, 0.95527864]], [[1.0, 1.0]], [[1.0, 1.0]]]]), + torch.ones((1, 2, 1, 2)), # data + torch.tensor([[2, 1, 0, 4], [-1, -3, 0, 5], [0, 0, 2.0, 5], [0, 0, 0, 1]]), + {}, + torch.tensor([[[[0.95527864, 0.95527864]], [[1.0, 1.0]], [[1.0, 1.0]]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": (3.0, 1.0), "padding_mode": "zeros"}, - np.arange(24).reshape((2, 3, 4)), # data - {"affine": np.diag([-3.0, 0.2, 1.5, 1])}, - np.array([[[0, 0], [4, 0], [8, 0]], [[12, 0], [16, 0], [20, 0]]]), + torch.arange(24).reshape((2, 3, 4)), # data + torch.as_tensor(np.diag([-3.0, 0.2, 1.5, 1])), + {}, + torch.tensor([[[0, 0], [4, 0], [8, 0]], [[12, 0], [16, 0], [20, 0]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": (3.0, 1.0), "padding_mode": "zeros"}, - np.arange(24).reshape((2, 3, 4)), # data + torch.arange(24).reshape((2, 3, 4)), # data + torch.eye(4), {}, - np.array([[[0, 1, 2, 3], [0, 0, 0, 0]], [[12, 13, 14, 15], [0, 0, 0, 0]]]), + torch.tensor([[[0, 1, 2, 3], [0, 0, 0, 0]], [[12, 13, 14, 15], [0, 0, 0, 0]]]), + *device, ] ) TESTS.append( [ - p, {"pixdim": (1.0, 1.0), "align_corners": True}, - np.arange(24).reshape((2, 3, 4)), # data + torch.arange(24).reshape((2, 3, 4)), # data + torch.eye(4), {}, - np.array( + torch.tensor( [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]] ), + *device, ] ) TESTS.append( [ - p, {"pixdim": (4.0, 5.0, 6.0)}, - np.arange(24).reshape((1, 2, 3, 4)), # data - {"affine": np.array([[-4, 0, 0, 4], [0, 5, 0, -5], [0, 0, 6, -6], [0, 0, 0, 1]])}, - np.arange(24).reshape((1, 2, 3, 4)), # data + torch.arange(24).reshape((1, 2, 3, 4)), # data + torch.tensor([[-4, 0, 0, 4], [0, 5, 0, -5], [0, 0, 6, -6], [0, 0, 0, 1]]), + {}, + torch.arange(24).reshape((1, 2, 3, 4)), # data + *device, ] ) TESTS.append( [ - p, {"pixdim": (4.0, 5.0, 6.0), "diagonal": True}, - np.arange(24).reshape((1, 2, 3, 4)), # data - {"affine": np.array([[-4, 0, 0, 4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]])}, - np.array( + torch.arange(24).reshape((1, 2, 3, 4)), # data + torch.tensor([[-4, 0, 0, 4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {}, + torch.tensor( [[[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]]] ), + *device, ] ) TESTS.append( [ - p, {"pixdim": (4.0, 5.0, 6.0), "padding_mode": "border", "diagonal": True}, - np.arange(24).reshape((1, 2, 3, 4)), # data - {"affine": np.array([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]])}, - np.array( + torch.arange(24).reshape((1, 2, 3, 4)), # data + torch.tensor([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {}, + torch.tensor( [[[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]]] ), + *device, ] ) TESTS.append( [ - p, {"pixdim": (4.0, 5.0, 6.0), "padding_mode": "border", "diagonal": True}, - np.arange(24).reshape((1, 2, 3, 4)), # data - {"affine": np.array([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), "mode": "nearest"}, - np.array( + torch.arange(24).reshape((1, 2, 3, 4)), # data + torch.tensor([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {"mode": "nearest"}, + torch.tensor( [[[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]]] ), + *device, ] ) TESTS.append( [ - p, {"pixdim": (1.9, 4.0), "padding_mode": "zeros", "diagonal": True}, - np.arange(24).reshape((1, 4, 6)), # data - {"affine": np.array([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), "mode": "nearest"}, - np.array( + torch.arange(24).reshape((1, 4, 6)), # data + torch.tensor([[-4, 0, 0, -4], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {"mode": "nearest"}, + torch.tensor( [ [ [18.0, 19.0, 20.0, 20.0, 21.0, 22.0, 23.0], @@ -148,15 +161,16 @@ ] ] ), + *device, ] ) TESTS.append( [ - p, - {"pixdim": (5.0, 3.0), "padding_mode": "border", "diagonal": True, "dtype": np.float32}, - np.arange(24).reshape((1, 4, 6)), # data - {"affine": np.array([[-4, 0, 0, 0], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), "mode": "bilinear"}, - np.array( + {"pixdim": (5.0, 3.0), "padding_mode": "border", "diagonal": True, "dtype": torch.float32}, + torch.arange(24).reshape((1, 4, 6)), # data + torch.tensor([[-4, 0, 0, 0], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {"mode": "bilinear"}, + torch.tensor( [ [ [18.0, 18.6, 19.2, 19.8, 20.400002, 21.0, 21.6, 22.2, 22.8], @@ -165,15 +179,16 @@ ] ] ), + *device, ] ) TESTS.append( [ - p, - {"pixdim": (5.0, 3.0), "padding_mode": "zeros", "diagonal": True, "dtype": np.float32}, - np.arange(24).reshape((1, 4, 6)), # data - {"affine": np.array([[-4, 0, 0, 0], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), "mode": "bilinear"}, - np.array( + {"pixdim": (5.0, 3.0), "padding_mode": "zeros", "diagonal": True, "dtype": torch.float32}, + torch.arange(24).reshape((1, 4, 6)), # data + torch.tensor([[-4, 0, 0, 0], [0, 5, 0, 0], [0, 0, 6, 0], [0, 0, 0, 1]]), + {"mode": "bilinear"}, + torch.tensor( [ [ [18.0000, 18.6000, 19.2000, 19.8000, 20.4000, 21.0000, 21.6000, 22.2000, 22.8000], @@ -182,45 +197,88 @@ ] ] ), + *device, ] ) TESTS.append( [ - p, {"pixdim": [-1, -1, 0.5], "padding_mode": "zeros", "dtype": float}, - np.ones((1, 2, 1, 2)), # data - {"affine": np.eye(4)}, - np.array([[[[1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0]]]]), + torch.ones((1, 2, 1, 2)), # data + torch.eye(4), + {}, + torch.tensor([[[[1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0]]]]), + *device, ] ) TESTS.append( # 5D input [ - p, {"pixdim": [-1, -1, 0.5], "padding_mode": "zeros", "dtype": float, "align_corners": True}, - np.ones((1, 2, 2, 2, 1)), # data - {"affine": np.eye(4)}, - np.ones((1, 2, 2, 3, 1)), + torch.ones((1, 2, 2, 2, 1)), # data + torch.eye(4), + {}, + torch.ones((1, 2, 2, 3, 1)), + *device, ] ) +TESTS_TORCH = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TESTS_TORCH.append([[1.2, 1.3, 0.9], torch.zeros((1, 3, 4, 5)), track_meta, *device]) + class TestSpacingCase(unittest.TestCase): @parameterized.expand(TESTS) - def test_spacing(self, in_type, init_param, img, data_param, expected_output): - _img = in_type(img) - output_data, _, new_affine = Spacing(**init_param)(_img, **data_param) - if isinstance(_img, torch.Tensor): - self.assertEqual(_img.device, output_data.device) - output_data = output_data.cpu() + def test_spacing(self, init_param, img, affine, data_param, expected_output, device): + img = MetaTensor(img, affine=affine).to(device) + res: MetaTensor = Spacing(**init_param)(img, **data_param) + self.assertEqual(img.device, res.device) - np.testing.assert_allclose(output_data, expected_output, atol=1e-1, rtol=1e-1) - sr = min(len(output_data.shape) - 1, 3) + assert_allclose(res, expected_output, atol=1e-1, rtol=1e-1) + sr = min(len(res.shape) - 1, 3) if isinstance(init_param["pixdim"], float): init_param["pixdim"] = [init_param["pixdim"]] * sr init_pixdim = ensure_tuple(init_param["pixdim"]) init_pixdim = init_param["pixdim"][:sr] - norm = affine_to_spacing(new_affine, sr) - np.testing.assert_allclose(fall_back_tuple(init_pixdim, norm), norm) + norm = affine_to_spacing(res.affine, sr).cpu().numpy() + assert_allclose(fall_back_tuple(init_pixdim, norm), norm, type_test=False) + + @parameterized.expand(TESTS_TORCH) + def test_spacing_torch(self, pixdim, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + tr = Spacing(pixdim=pixdim) + img = img.to(device) + res = tr(img) + if track_meta: + self.assertIsInstance(res, MetaTensor) + new_spacing = affine_to_spacing(res.affine, 3) + assert_allclose(new_spacing, pixdim, type_test=False) + self.assertNotEqual(img.shape, res.shape) + else: + self.assertIsInstance(res, torch.Tensor) + self.assertNotIsInstance(res, MetaTensor) + self.assertNotEqual(img.shape, res.shape) + + @parameterized.expand(TEST_DEVICES) + def test_inverse(self, device): + img_t = torch.rand((1, 10, 9, 8), dtype=torch.float32, device=device) + affine = torch.tensor( + [[0, 0, -1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=torch.float32, device=device + ) + meta = {"fname": "somewhere"} + img = MetaTensor(img_t, affine=affine, meta=meta) + tr = Spacing(pixdim=[1.1, 1.2, 0.9]) + # check that image and affine have changed + img = tr(img) + self.assertNotEqual(img.shape, img_t.shape) + l2_norm_affine = ((affine - img.affine) ** 2).sum() ** 0.5 + self.assertGreater(l2_norm_affine, 5e-2) + # check that with inverse, image affine are back to how they were + img = tr.inverse(img) + self.assertEqual(img.applied_operations, []) + self.assertEqual(img.shape, img_t.shape) + l2_norm_affine = ((affine - img.affine) ** 2).sum() ** 0.5 + self.assertLess(l2_norm_affine, 5e-2) if __name__ == "__main__": diff --git a/tests/test_spacingd.py b/tests/test_spacingd.py index 060d908699..d3c7bbc629 100644 --- a/tests/test_spacingd.py +++ b/tests/test_spacingd.py @@ -16,83 +16,107 @@ import torch from parameterized import parameterized +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor +from monai.data.utils import affine_to_spacing from monai.transforms import Spacingd -from monai.utils.enums import PostFix -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_DEVICES, assert_allclose TESTS: List[Tuple] = [] -for p in TEST_NDARRAYS: +for device in TEST_DEVICES: TESTS.append( ( "spacing 3d", - {"image": p(np.ones((2, 10, 15, 20))), PostFix.meta("image"): {"affine": p(np.eye(4))}}, + {"image": MetaTensor(torch.ones((2, 10, 15, 20)), affine=torch.eye(4))}, dict(keys="image", pixdim=(1, 2, 1.4)), - ("image", PostFix.meta("image"), "image_transforms"), (2, 10, 8, 15), - p(np.diag([1, 2, 1.4, 1.0])), + torch.as_tensor(np.diag([1, 2, 1.4, 1.0])), + *device, ) ) TESTS.append( ( "spacing 2d", - {"image": np.ones((2, 10, 20)), PostFix.meta("image"): {"affine": np.eye(3)}}, + {"image": MetaTensor(torch.ones((2, 10, 20)), affine=torch.eye(3))}, dict(keys="image", pixdim=(1, 2)), - ("image", PostFix.meta("image"), "image_transforms"), (2, 10, 10), - np.diag((1, 2, 1)), + torch.as_tensor(np.diag((1, 2, 1))), + *device, ) ) TESTS.append( ( "spacing 2d no metadata", - {"image": np.ones((2, 10, 20))}, + {"image": MetaTensor(torch.ones((2, 10, 20)))}, dict(keys="image", pixdim=(1, 2)), - ("image", PostFix.meta("image"), "image_transforms"), (2, 10, 10), - np.diag((1, 2, 1)), + torch.as_tensor(np.diag((1, 2, 1))), + *device, ) ) TESTS.append( ( "interp all", { - "image": np.arange(20).reshape((2, 1, 10)), - "seg": np.ones((2, 1, 10)), - PostFix.meta("image"): {"affine": np.eye(4)}, - PostFix.meta("seg"): {"affine": np.eye(4)}, + "image": MetaTensor(np.arange(20).reshape((2, 1, 10)), affine=torch.eye(4)), + "seg": MetaTensor(torch.ones((2, 1, 10)), affine=torch.eye(4)), }, dict(keys=("image", "seg"), mode="nearest", pixdim=(1, 0.2)), - ("image", PostFix.meta("image"), "image_transforms", "seg", PostFix.meta("seg"), "seg_transforms"), (2, 1, 46), - np.diag((1, 0.2, 1, 1)), + torch.as_tensor(np.diag((1, 0.2, 1))), + *device, ) ) TESTS.append( ( "interp sep", { - "image": np.ones((2, 1, 10)), - "seg": np.ones((2, 1, 10)), - PostFix.meta("image"): {"affine": np.eye(4)}, - PostFix.meta("seg"): {"affine": np.eye(4)}, + "image": MetaTensor(torch.ones((2, 1, 10)), affine=torch.eye(4)), + "seg": MetaTensor(torch.ones((2, 1, 10)), affine=torch.eye(4)), }, dict(keys=("image", "seg"), mode=("bilinear", "nearest"), pixdim=(1, 0.2)), - ("image", PostFix.meta("image"), "image_transforms", "seg", PostFix.meta("seg"), "seg_transforms"), (2, 1, 46), - np.diag((1, 0.2, 1, 1)), + torch.as_tensor(np.diag((1, 0.2, 1))), + *device, ) ) +TESTS_TORCH = [] +for track_meta in (False, True): + for device in TEST_DEVICES: + TESTS_TORCH.append([{"keys": "seg", "pixdim": [0.2, 0.3, 1]}, torch.ones(2, 1, 2, 3), track_meta, *device]) + + class TestSpacingDCase(unittest.TestCase): @parameterized.expand(TESTS) - def test_spacingd(self, _, data, kw_args, expected_keys, expected_shape, expected_affine): + def test_spacingd(self, _, data, kw_args, expected_shape, expected_affine, device): + data = {k: v.to(device) for k, v in data.items()} res = Spacingd(**kw_args)(data) - if isinstance(data["image"], torch.Tensor): - self.assertEqual(data["image"].device, res["image"].device) - self.assertEqual(expected_keys, tuple(sorted(res))) - np.testing.assert_allclose(res["image"].shape, expected_shape) - assert_allclose(res[PostFix.meta("image")]["affine"], expected_affine) + in_img = data["image"] + out_img = res["image"] + self.assertEqual(in_img.device, out_img.device) + # no change in number of keys + self.assertEqual(tuple(sorted(data)), tuple(sorted(res))) + np.testing.assert_allclose(out_img.shape, expected_shape) + assert_allclose(out_img.affine, expected_affine) + + @parameterized.expand(TESTS_TORCH) + def test_orntd_torch(self, init_param, img: torch.Tensor, track_meta: bool, device): + set_track_meta(track_meta) + tr = Spacingd(**init_param) + data = {"seg": img.to(device)} + res = tr(data)["seg"] + + if track_meta: + self.assertIsInstance(res, MetaTensor) + new_spacing = affine_to_spacing(res.affine, 3) + assert_allclose(new_spacing, init_param["pixdim"], type_test=False) + self.assertNotEqual(img.shape, res.shape) + else: + self.assertIsInstance(res, torch.Tensor) + self.assertNotIsInstance(res, MetaTensor) + self.assertNotEqual(img.shape, res.shape) if __name__ == "__main__": diff --git a/tests/test_spatial_crop.py b/tests/test_spatial_crop.py index bf1eb11491..6fdfbd3f70 100644 --- a/tests/test_spatial_crop.py +++ b/tests/test_spatial_crop.py @@ -11,12 +11,10 @@ import unittest -import numpy as np -import torch from parameterized import parameterized from monai.transforms import SpatialCrop -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.croppers import CropTest TESTS = [ [{"roi_center": [1, 1, 1], "roi_size": [2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2)], @@ -26,31 +24,28 @@ [{"roi_start": [0, 0, 0, 0, 0], "roi_end": [2, 2, 2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2)], [{"roi_start": [0, 0, 0, 0, 0], "roi_end": [8, 8, 8, 2, 2]}, (3, 3, 3, 3), (3, 3, 3, 3)], [{"roi_start": [1, 0, 0], "roi_end": [1, 8, 8]}, (3, 3, 3, 3), (3, 0, 3, 3)], - [{"roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]}, (3, 3, 3, 3), (3, 1, 2, 2)], + [ + {"roi_slices": [slice(s, e) for s, e in zip([None, None, None], [None, None, None])]}, + (3, 11, 12, 15), + (3, 11, 12, 15), + ], + [{"roi_slices": [slice(s, e) for s, e in zip([1, None, 0], [None, None, None])]}, (3, 7, 9, 11), (3, 6, 9, 11)], + [{"roi_slices": [slice(s, e) for s, e in zip([0, None, None], [-1, None, None])]}, (3, 7, 9, 11), (3, 6, 9, 11)], + [{"roi_slices": [slice(s, e) for s, e in zip([1, None, None], [None, None, None])]}, (3, 10, 8, 6), (3, 9, 8, 6)], + [{"roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]}, (3, 15, 17, 8), (3, 1, 2, 2)], + [{"roi_slices": [slice(s, e) for s, e in zip([None, None, None], [-2, -1, 2])]}, (3, 13, 8, 6), (3, 11, 7, 2)], + [{"roi_start": [-1, 0], "roi_end": [5, 5]}, (1, 5, 5), (1, 5, 5)], ] TEST_ERRORS = [[{"roi_slices": [slice(s, e, 2) for s, e in zip([-1, -2, 0], [None, None, 2])]}]] -class TestSpatialCrop(unittest.TestCase): +class TestSpatialCrop(CropTest): + Cropper = SpatialCrop + @parameterized.expand(TESTS) def test_shape(self, input_param, input_shape, expected_shape): - input_data = np.random.randint(0, 2, size=input_shape) - results = [] - for p in TEST_NDARRAYS: - for q in TEST_NDARRAYS + (None,): - input_param_mod = { - k: q(v) if k != "roi_slices" and q is not None else v for k, v in input_param.items() - } - im = p(input_data) - result = SpatialCrop(**input_param_mod)(im) - self.assertEqual(type(im), type(result)) - if isinstance(result, torch.Tensor): - self.assertEqual(result.device, im.device) - self.assertTupleEqual(result.shape, expected_shape) - results.append(result) - if len(results) > 1: - assert_allclose(results[0], results[-1], type_test=False) + self.crop_test(input_param, input_shape, expected_shape) @parameterized.expand(TEST_ERRORS) def test_error(self, input_param): diff --git a/tests/test_spatial_cropd.py b/tests/test_spatial_cropd.py index 5b16f460fd..11f6da0811 100644 --- a/tests/test_spatial_cropd.py +++ b/tests/test_spatial_cropd.py @@ -11,56 +11,57 @@ import unittest -import numpy as np from parameterized import parameterized from monai.transforms import SpatialCropd -from tests.utils import TEST_NDARRAYS +from tests.croppers import CropTest -TESTS = [] -for p in TEST_NDARRAYS: - TESTS.append( - [ - {"keys": ["img"], "roi_center": [1, 1, 1], "roi_size": [2, 2, 2]}, - {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, - (3, 2, 2, 2), - ] - ) - TESTS.append( - [ - {"keys": ["img"], "roi_start": [0, 0, 0], "roi_end": [2, 2, 2]}, - {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, - (3, 2, 2, 2), - ] - ) - TESTS.append( - [ - {"keys": ["img"], "roi_start": [0, 0], "roi_end": [2, 2]}, - {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, - (3, 2, 2, 3), - ] - ) - TESTS.append( - [ - {"keys": ["img"], "roi_start": [0, 0, 0, 0, 0], "roi_end": [2, 2, 2, 2, 2]}, - {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, - (3, 2, 2, 2), - ] - ) - TESTS.append( - [ - {"keys": ["img"], "roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]}, - {"img": p(np.random.randint(0, 2, size=[3, 3, 3, 3]))}, - (3, 1, 2, 2), - ] - ) +TESTS = [ + [ + {"keys": ["img"], "roi_center": [1, 1], "roi_size": [2, 2]}, + (1, 3, 3), + (1, 2, 2), + (slice(None), slice(None, 2), slice(None, 2)), + ], + [ + {"keys": ["img"], "roi_center": [1, 1, 1], "roi_size": [2, 2, 2]}, + (3, 3, 3, 3), + (3, 2, 2, 2), + (slice(None), slice(None, 2), slice(None, 2), slice(None, 2)), + ], + [ + {"keys": ["img"], "roi_start": [0, 0, 0], "roi_end": [2, 2, 2]}, + (3, 3, 3, 3), + (3, 2, 2, 2), + (slice(None), slice(None, 2), slice(None, 2), slice(None, 2)), + ], + [ + {"keys": ["img"], "roi_start": [0, 0], "roi_end": [2, 2]}, + (3, 3, 3, 3), + (3, 2, 2, 3), + (slice(None), slice(None, 2), slice(None, 2), slice(None)), + ], + [ + {"keys": ["img"], "roi_start": [0, 0, 0, 0, 0], "roi_end": [2, 2, 2, 2, 2]}, + (3, 3, 3, 3), + (3, 2, 2, 2), + (slice(None), slice(None, 2), slice(None, 2), slice(None, 2)), + ], + [ + {"keys": ["img"], "roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]}, + (3, 3, 3, 3), + (3, 1, 2, 2), + (slice(None), slice(-1, None), slice(-2, None), slice(0, 2)), + ], +] -class TestSpatialCropd(unittest.TestCase): +class TestSpatialCropd(CropTest): + Cropper = SpatialCropd + @parameterized.expand(TESTS) - def test_shape(self, input_param, input_data, expected_shape): - result = SpatialCropd(**input_param)(input_data) - self.assertTupleEqual(result["img"].shape, expected_shape) + def test_shape(self, input_param, input_shape, expected_shape, same_area): + self.crop_test(input_param, input_shape, expected_shape, same_area) if __name__ == "__main__": diff --git a/tests/test_spatial_pad.py b/tests/test_spatial_pad.py index 4cdeb6d64e..5a70c10686 100644 --- a/tests/test_spatial_pad.py +++ b/tests/test_spatial_pad.py @@ -10,91 +10,28 @@ # limitations under the License. import unittest -from typing import List -import numpy as np -import torch from parameterized import parameterized from monai.transforms import SpatialPad -from monai.utils.enums import NumpyPadMode, PytorchPadMode -from monai.utils.misc import set_determinism -from tests.utils import TEST_NDARRAYS +from tests.padders import PadTest TESTS = [] +TESTS.append([{"spatial_size": [3, 4], "method": "end"}, (1, 2, 3), (1, 3, 4)]) +TESTS.append([{"spatial_size": [15, 4, -1], "method": "symmetric"}, (3, 8, 8, 4), (3, 15, 8, 4)]) -MODES = [] -# Test modes -NP_MODES: List = [ - "constant", - "edge", - # `reflect` mode is not supported in some PyTorch versions, skip the test - # "reflect", - "wrap", -] -MODES += NP_MODES -MODES += [NumpyPadMode(i) for i in NP_MODES] - -PT_MODES: list = [ - "constant", - "replicate", - "circular", - # `reflect` mode is not supported in some PyTorch versions, skip the test - # "reflect", -] -MODES += PT_MODES -MODES += [PytorchPadMode(i) for i in PT_MODES] - -for mode in MODES: - TESTS.append([{"spatial_size": [3, 4], "method": "end", "mode": mode}, (1, 2, 3), (1, 3, 4)]) - - TESTS.append([{"spatial_size": [15, 4, -1], "method": "symmetric", "mode": mode}, (3, 8, 8, 4), (3, 15, 8, 4)]) - - -class TestSpatialPad(unittest.TestCase): - def setUp(self) -> None: - set_determinism(seed=0) - - def tearDown(self) -> None: - set_determinism(None) - - @staticmethod - def get_arr(shape): - return np.random.randint(100, size=shape).astype(float) +class TestSpatialPad(PadTest): + Padder = SpatialPad @parameterized.expand(TESTS) - def test_pad_shape(self, input_param, input_shape, expected_shape): - results_1 = [] - results_2 = [] - input_data = self.get_arr(input_shape) - # check result is the same regardless of input type - for p in TEST_NDARRAYS: - padder = SpatialPad(**input_param) - r1 = padder(p(input_data)) - r2 = padder(p(input_data), mode=input_param["mode"]) - results_1.append(r1.cpu() if isinstance(r1, torch.Tensor) else r1) - results_2.append(r2.cpu() if isinstance(r2, torch.Tensor) else r2) - for results in (results_1, results_2): - np.testing.assert_allclose(results[-1].shape, expected_shape) - if input_param["mode"] not in ("empty", NumpyPadMode.EMPTY): - torch.testing.assert_allclose(results[0], results[-1], atol=0, rtol=1e-5) + def test_pad(self, input_param, input_shape, expected_shape): + self.pad_test(input_param, input_shape, expected_shape) def test_pad_kwargs(self): - for p in TEST_NDARRAYS: - input_data = p(np.zeros((3, 8, 4))) - if isinstance(input_data, torch.Tensor): - result = ( - SpatialPad(spatial_size=[15, 8], method="end", mode="constant", value=2)(img=input_data) - .cpu() - .numpy() - ) - else: - result = SpatialPad( - spatial_size=[15, 8], method="end", mode="constant", constant_values=((0, 0), (1, 1), (2, 2)) - )(img=input_data) - torch.testing.assert_allclose(result[:, 8:, :4], np.ones((3, 7, 4)), rtol=1e-7, atol=0) - torch.testing.assert_allclose(result[:, :, 4:], np.ones((3, 15, 4)) + 1, rtol=1e-7, atol=0) + kwargs = {"spatial_size": [15, 8], "method": "end", "mode": "constant"} + unchanged_slices = [slice(None), slice(None, 8), slice(None, 4)] + self.pad_test_kwargs(unchanged_slices, **kwargs) if __name__ == "__main__": diff --git a/tests/test_spatial_padd.py b/tests/test_spatial_padd.py index 762a1145f5..656a731de0 100644 --- a/tests/test_spatial_padd.py +++ b/tests/test_spatial_padd.py @@ -11,42 +11,26 @@ import unittest -import numpy as np from parameterized import parameterized from monai.transforms import SpatialPadd +from tests.padders import PadTest -TEST_CASE_1 = [ - {"keys": ["img"], "spatial_size": [15, 8, 8], "method": "symmetric", "mode": "constant"}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 15, 8, 8)), +TESTS = [ + [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "symmetric"}, (3, 8, 8, 4), (3, 15, 8, 8)], + [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end"}, (3, 8, 8, 4), (3, 15, 8, 8)], + [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end"}, (3, 8, 8, 4), (3, 15, 8, 8)], + [{"keys": ["img"], "spatial_size": [15, 8, -1], "method": "end"}, (3, 8, 4, 4), (3, 15, 8, 4)], ] -TEST_CASE_2 = [ - {"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end", "mode": "constant"}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 15, 8, 8)), -] - -TEST_CASE_3 = [ - {"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end", "mode": {"constant"}}, - {"img": np.zeros((3, 8, 8, 4))}, - np.zeros((3, 15, 8, 8)), -] - -TEST_CASE_4 = [ - {"keys": ["img"], "spatial_size": [15, 8, -1], "method": "end", "mode": {"constant"}}, - {"img": np.zeros((3, 8, 4, 4))}, - np.zeros((3, 15, 8, 4)), -] +class TestSpatialPadd(PadTest): + Padder = SpatialPadd -class TestSpatialPadd(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) - def test_pad_shape(self, input_param, input_data, expected_val): - padder = SpatialPadd(**input_param) - result = padder(input_data) - np.testing.assert_allclose(result["img"].shape, expected_val.shape) + @parameterized.expand(TESTS) + def test_pad(self, input_param, input_shape, expected_shape): + modes = ["constant", {"constant"}] + self.pad_test(input_param, input_shape, expected_shape, modes) if __name__ == "__main__": diff --git a/tests/test_spatial_resample.py b/tests/test_spatial_resample.py index 9ee84de85b..63260373d0 100644 --- a/tests/test_spatial_resample.py +++ b/tests/test_spatial_resample.py @@ -9,137 +9,196 @@ # See the License for the specific language governing permissions and # limitations under the License. -import itertools import unittest import numpy as np +import torch from parameterized import parameterized from monai.config import USE_COMPILED +from monai.data.meta_obj import set_track_meta +from monai.data.meta_tensor import MetaTensor +from monai.data.utils import to_affine_nd from monai.transforms import SpatialResample -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_DEVICES, TEST_NDARRAYS_ALL, assert_allclose TESTS = [] -for ind, dst in enumerate( - [ - np.asarray([[1.0, 0.0, 0.0], [0.0, -1.0, 1.0], [0.0, 0.0, 1.0]]), # flip the second - np.asarray([[-1.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), # flip the first - ] -): - for p in TEST_NDARRAYS: - for p_src in TEST_NDARRAYS: - for align in (False, True): - for interp_mode in ("nearest", "bilinear"): - TESTS.append( - [ - {}, # default no params - np.arange(4).reshape((1, 2, 2)) + 1.0, # data - { - "src_affine": p_src(np.eye(3)), - "dst_affine": p(dst), - "dtype": np.float32, - "align_corners": align, - "mode": interp_mode, - "padding_mode": "zeros", - }, - np.array([[[2.0, 1.0], [4.0, 3.0]]]) if ind == 0 else np.array([[[3.0, 4.0], [1.0, 2.0]]]), - ] - ) -for ind, dst in enumerate( - [ - np.asarray([[1.0, 0.0, 0.0, 0.0], [0.0, -1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), - np.asarray([[-1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), - ] -): - for p_src in TEST_NDARRAYS: - for align in (True, False): - if align and USE_COMPILED: - interp = ("nearest", "bilinear", 0, 1) - else: - interp = ("nearest", "bilinear") # type: ignore - for interp_mode in interp: # type: ignore +destinations_3d = [ + torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, -1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), + torch.tensor([[-1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), +] +expected_3d = [ + torch.tensor([[[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[10.0, 11.0, 12.0], [7.0, 8.0, 9.0]]]]), + torch.tensor([[[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]]), +] + +for dst, expct in zip(destinations_3d, expected_3d): + for device in TEST_DEVICES: + for align in (False, True): + interp = ("nearest", "bilinear", 0, 1) if align and USE_COMPILED else ("nearest", "bilinear") + for interp_mode in interp: for padding_mode in ("zeros", "border", "reflection"): TESTS.append( [ - {}, # default no params - np.arange(12).reshape((1, 2, 2, 3)) + 1.0, # data + torch.arange(12).reshape((1, 2, 2, 3)) + 1.0, # data + *device, { - "src_affine": p_src(np.eye(4)), - "dst_affine": p_src(dst), - "dtype": np.float64, + "dst_affine": dst, + "dtype": torch.float64, "align_corners": align, "mode": interp_mode, "padding_mode": padding_mode, }, - np.array([[[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[10.0, 11.0, 12.0], [7.0, 8.0, 9.0]]]]) - if ind == 0 - else np.array( - [[[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]] - ), + expct, ] ) -class TestSpatialResample(unittest.TestCase): - @parameterized.expand(itertools.product(TEST_NDARRAYS, TESTS)) - def test_flips(self, p_type, args): - init_param, img, data_param, expected_output = args - _img = p_type(img) - _expected_output = p_type(expected_output) - output_data, output_dst = SpatialResample(**init_param)(img=_img, **data_param) - assert_allclose(output_data, _expected_output, rtol=1e-2, atol=1e-2) - expected_dst = ( - data_param.get("dst_affine") if data_param.get("dst_affine") is not None else data_param.get("src_affine") - ) - assert_allclose(output_dst, expected_dst, type_test=False, rtol=1e-2, atol=1e-2) +destinations_2d = [ + torch.tensor([[1.0, 0.0, 0.0], [0.0, -1.0, 1.0], [0.0, 0.0, 1.0]]), # flip the second + torch.tensor([[-1.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), # flip the first +] +expected_2d = [torch.tensor([[[2.0, 1.0], [4.0, 3.0]]]), torch.tensor([[[3.0, 4.0], [1.0, 2.0]]])] - @parameterized.expand(itertools.product([True, False], TEST_NDARRAYS)) - def test_4d_5d(self, is_5d, p_type): - new_shape = (1, 2, 2, 3, 1, 1) if is_5d else (1, 2, 2, 3, 1) - img = np.arange(12).reshape(new_shape) - img = np.tile(img, (1, 1, 1, 1, 2, 2) if is_5d else (1, 1, 1, 1, 2)) - _img = p_type(img) - dst = np.asarray([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 1.5], [0.0, 0.0, 0.0, 1.0]]) - output_data, output_dst = SpatialResample(dtype=np.float32)( - img=_img, src_affine=p_type(np.eye(4)), dst_affine=dst - ) - expected_data = ( - np.asarray( - [ +for dst, expct in zip(destinations_2d, expected_2d): + for device in TEST_DEVICES: + for align in (False, True): + for interp_mode in ("nearest", "bilinear"): + TESTS.append( [ - [[[0.0, 0.0], [0.0, 1.0]], [[0.5, 0.0], [1.5, 1.0]], [[1.0, 2.0], [2.0, 2.0]]], - [[[3.0, 3.0], [3.0, 4.0]], [[3.5, 3.0], [4.5, 4.0]], [[4.0, 5.0], [5.0, 5.0]]], - ], + torch.arange(4).reshape((1, 2, 2)) + 1.0, + *device, + { + "dst_affine": dst, + "dtype": torch.float32, + "align_corners": align, + "mode": interp_mode, + "padding_mode": "zeros", + }, + expct, + ] + ) + +TEST_4_5_D = [] +for device in TEST_DEVICES: + for dtype in (torch.float32, torch.float64): + # 4D + TEST_4_5_D.append( + [ + (1, 2, 2, 3, 1), + (1, 1, 1, 1, 2), + *device, + dtype, + torch.tensor( [ - [[[6.0, 6.0], [6.0, 7.0]], [[6.5, 6.0], [7.5, 7.0]], [[7.0, 8.0], [8.0, 8.0]]], - [[[9.0, 9.0], [9.0, 10.0]], [[9.5, 9.0], [10.5, 10.0]], [[10.0, 11.0], [11.0, 11.0]]], - ], - ], - dtype=np.float32, - ) - if is_5d - else np.asarray( - [ - [[[0.5, 0.0], [0.0, 2.0], [1.5, 1.0]], [[3.5, 3.0], [3.0, 5.0], [4.5, 4.0]]], - [[[6.5, 6.0], [6.0, 8.0], [7.5, 7.0]], [[9.5, 9.0], [9.0, 11.0], [10.5, 10.0]]], - ], - dtype=np.float32, - ) + [[[0.5, 0.0], [0.0, 2.0], [1.5, 1.0]], [[3.5, 3.0], [3.0, 5.0], [4.5, 4.0]]], + [[[6.5, 6.0], [6.0, 8.0], [7.5, 7.0]], [[9.5, 9.0], [9.0, 11.0], [10.5, 10.0]]], + ] + ), + ] ) - assert_allclose(output_data, p_type(expected_data[None]), rtol=1e-2, atol=1e-2) - assert_allclose(output_dst, dst, type_test=False, rtol=1e-2, atol=1e-2) - - def test_ill_affine(self): - img = np.arange(12).reshape(1, 2, 2, 3) - ill_affine = np.asarray( - [[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, -1.0, 1.5], [0.0, 0.0, 0.0, 1.0]] + # 5D + TEST_4_5_D.append( + [ + (1, 2, 2, 3, 1, 1), + (1, 1, 1, 1, 2, 2), + *device, + dtype, + torch.tensor( + [ + [ + [[[0.0, 0.0], [0.0, 1.0]], [[0.5, 0.0], [1.5, 1.0]], [[1.0, 2.0], [2.0, 2.0]]], + [[[3.0, 3.0], [3.0, 4.0]], [[3.5, 3.0], [4.5, 4.0]], [[4.0, 5.0], [5.0, 5.0]]], + ], + [ + [[[6.0, 6.0], [6.0, 7.0]], [[6.5, 6.0], [7.5, 7.0]], [[7.0, 8.0], [8.0, 8.0]]], + [[[9.0, 9.0], [9.0, 10.0]], [[9.5, 9.0], [10.5, 10.0]], [[10.0, 11.0], [11.0, 11.0]]], + ], + ] + ), + ] ) + +TEST_TORCH_INPUT = [] +for track_meta in (True,): + for t in TEST_4_5_D: + TEST_TORCH_INPUT.append(t + [track_meta]) + + +class TestSpatialResample(unittest.TestCase): + @parameterized.expand(TESTS) + def test_flips(self, img, device, data_param, expected_output): + for p in TEST_NDARRAYS_ALL: + img = p(img) + if isinstance(img, MetaTensor): + img.affine = torch.eye(4) + if hasattr(img, "to"): + img = img.to(device) + out = SpatialResample()(img=img, **data_param) + assert_allclose(out, expected_output, rtol=1e-2, atol=1e-2) + assert_allclose(out.affine, data_param["dst_affine"]) + + @parameterized.expand(TEST_4_5_D) + def test_4d_5d(self, new_shape, tile, device, dtype, expected_data): + img = np.arange(12).reshape(new_shape) + img = np.tile(img, tile) + img = MetaTensor(img).to(device) + + dst = torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 1.5], [0.0, 0.0, 0.0, 1.0]]) + dst = dst.to(dtype) + out = SpatialResample(dtype=dtype)(img=img, dst_affine=dst) + assert_allclose(out, expected_data[None], rtol=1e-2, atol=1e-2) + assert_allclose(out.affine, dst.to(torch.float32), rtol=1e-2, atol=1e-2) + + @parameterized.expand(TEST_DEVICES) + def test_ill_affine(self, device): + img = MetaTensor(torch.arange(12).reshape(1, 2, 2, 3)).to(device) + ill_affine = torch.tensor([[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, -1, 1.5], [0, 0, 0, 1]]) with self.assertRaises(ValueError): - SpatialResample()(img=img, src_affine=np.eye(4), dst_affine=ill_affine) + img.affine = torch.eye(4) + dst_affine = ill_affine + SpatialResample()(img=img, dst_affine=dst_affine) with self.assertRaises(ValueError): - SpatialResample()(img=img, src_affine=ill_affine, dst_affine=np.eye(3)) + img.affine = ill_affine + dst_affine = torch.eye(4) + SpatialResample()(img=img, dst_affine=dst_affine) + with self.assertRaises(ValueError): + img.affine = torch.eye(4) + dst_affine = torch.eye(4) * 0.1 + SpatialResample(mode=None)(img=img, dst_affine=dst_affine) + + @parameterized.expand(TEST_TORCH_INPUT) + def test_input_torch(self, new_shape, tile, device, dtype, expected_data, track_meta): + set_track_meta(track_meta) + img = np.arange(12).reshape(new_shape) + img = torch.as_tensor(np.tile(img, tile)).to(device) + dst = torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 1.5], [0.0, 0.0, 0.0, 1.0]]) + dst = dst.to(dtype).to(device) + + out = SpatialResample(dtype=dtype)(img=img, dst_affine=dst) + assert_allclose(out, expected_data[None], rtol=1e-2, atol=1e-2) + if track_meta: + self.assertIsInstance(out, MetaTensor) + assert_allclose(out.affine, dst.to(torch.float32), rtol=1e-2, atol=1e-2) + else: + self.assertIsInstance(out, torch.Tensor) + self.assertNotIsInstance(out, MetaTensor) + + @parameterized.expand(TESTS) + def test_inverse(self, img, device, data_param, expected_output): + img = MetaTensor(img, affine=torch.eye(4)).to(device) + tr = SpatialResample() + out = tr(img=img, **data_param) + assert_allclose(out, expected_output, rtol=1e-2, atol=1e-2) + assert_allclose(out.affine, data_param["dst_affine"]) + + # inverse + out = tr.inverse(out) + assert_allclose(img, out) + expected_affine = to_affine_nd(len(out.affine) - 1, torch.eye(4)) + assert_allclose(out.affine, expected_affine) if __name__ == "__main__": diff --git a/tests/test_spatial_resampled.py b/tests/test_spatial_resampled.py index 73f83791d9..3772cf0ddf 100644 --- a/tests/test_spatial_resampled.py +++ b/tests/test_spatial_resampled.py @@ -9,104 +9,100 @@ # See the License for the specific language governing permissions and # limitations under the License. -import itertools import unittest import numpy as np +import torch from parameterized import parameterized from monai.config import USE_COMPILED -from monai.transforms import SpatialResampleD -from tests.utils import TEST_NDARRAYS, assert_allclose +from monai.data.meta_tensor import MetaTensor +from monai.data.utils import to_affine_nd +from monai.transforms.spatial.dictionary import SpatialResampled +from tests.utils import TEST_DEVICES, assert_allclose TESTS = [] -for ind, dst in enumerate( - [ - np.asarray([[1.0, 0.0, 0.0], [0.0, -1.0, 1.0], [0.0, 0.0, 1.0]]), # flip the second - np.asarray([[-1.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), # flip the first - ] -): - for p in TEST_NDARRAYS: - for p_src in TEST_NDARRAYS: - for align in (False, True): + +destinations_3d = [ + torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, -1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), + torch.tensor([[-1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), +] +expected_3d = [ + torch.tensor([[[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[10.0, 11.0, 12.0], [7.0, 8.0, 9.0]]]]), + torch.tensor([[[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]]), +] + +for dst, expct in zip(destinations_3d, expected_3d): + for device in TEST_DEVICES: + for align in (True, False): + for dtype in (torch.float32, torch.float64): + interp = ("nearest", "bilinear", 0, 1) if align and USE_COMPILED else ("nearest", "bilinear") + for interp_mode in interp: + for padding_mode in ("zeros", "border", "reflection"): + TESTS.append( + [ + np.arange(12).reshape((1, 2, 2, 3)) + 1.0, # data + *device, + dst, + { + "dst_keys": "dst_affine", + "dtype": dtype, + "align_corners": align, + "mode": interp_mode, + "padding_mode": padding_mode, + }, + expct, + ] + ) + +destinations_2d = [ + torch.tensor([[1.0, 0.0, 0.0], [0.0, -1.0, 1.0], [0.0, 0.0, 1.0]]), # flip the second + torch.tensor([[-1.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), # flip the first +] + +expected_2d = [torch.tensor([[[2.0, 1.0], [4.0, 3.0]]]), torch.tensor([[[3.0, 4.0], [1.0, 2.0]]])] + +for dst, expct in zip(destinations_2d, expected_2d): + for device in TEST_DEVICES: + for align in (False, True): + for dtype in (torch.float32, torch.float64): for interp_mode in ("nearest", "bilinear"): TESTS.append( [ - {}, # default no params np.arange(4).reshape((1, 2, 2)) + 1.0, # data + *device, + dst, { - "src": p_src(np.eye(3)), - "dst": p(dst), - "dtype": np.float32, + "dst_keys": "dst_affine", + "dtype": dtype, "align_corners": align, "mode": interp_mode, "padding_mode": "zeros", }, - np.array([[[2.0, 1.0], [4.0, 3.0]]]) if ind == 0 else np.array([[[3.0, 4.0], [1.0, 2.0]]]), - ] - ) - -for ind, dst in enumerate( - [ - np.asarray([[1.0, 0.0, 0.0, 0.0], [0.0, -1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), - np.asarray([[-1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]), - ] -): - for p_src in TEST_NDARRAYS: - for align in (True, False): - if align and USE_COMPILED: - interp = ("nearest", "bilinear", 0, 1) - else: - interp = ("nearest", "bilinear") # type: ignore - for interp_mode in interp: # type: ignore - for padding_mode in ("zeros", "border", "reflection"): - TESTS.append( - [ - {}, # default no params - np.arange(12).reshape((1, 2, 2, 3)) + 1.0, # data - { - "src": p_src(np.eye(4)), - "dst": p_src(dst), - "dtype": np.float64, - "align_corners": align, - "mode": interp_mode, - "padding_mode": padding_mode, - }, - np.array([[[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[10.0, 11.0, 12.0], [7.0, 8.0, 9.0]]]]) - if ind == 0 - else np.array( - [[[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]] - ), + expct, ] ) class TestSpatialResample(unittest.TestCase): - @parameterized.expand(itertools.product(TEST_NDARRAYS, TESTS)) - def test_flips_inverse(self, p_type, args): - _, img, data_param, expected_output = args - _img = p_type(img) - _expected_output = p_type(expected_output) - input_dict = {"img": _img, "img_meta_dict": {"src": data_param.get("src"), "dst": data_param.get("dst")}} - xform = SpatialResampleD( - keys="img", - meta_src_keys="src", - meta_dst_keys="dst", - mode=data_param.get("mode"), - padding_mode=data_param.get("padding_mode"), - align_corners=data_param.get("align_corners"), - ) - output_data = xform(input_dict) - assert_allclose(output_data["img"], _expected_output, rtol=1e-2, atol=1e-2) - assert_allclose( - output_data["img_meta_dict"]["src"], data_param.get("dst"), type_test=False, rtol=1e-2, atol=1e-2 - ) - - inverted = xform.inverse(output_data) - self.assertEqual(inverted["img_transforms"], []) # no further invert after inverting - assert_allclose(inverted["img_meta_dict"]["src"], data_param.get("src"), type_test=False, rtol=1e-2, atol=1e-2) - assert_allclose(inverted["img"], _img, rtol=1e-2, atol=1e-2) + @parameterized.expand(TESTS) + def test_flips_inverse(self, img, device, dst_affine, kwargs, expected_output): + img = MetaTensor(img, affine=torch.eye(4)).to(device) + data = {"img": img, "dst_affine": dst_affine} + + xform = SpatialResampled(keys="img", **kwargs) + output_data = xform(data) + out = output_data["img"] + + assert_allclose(out, expected_output, rtol=1e-2, atol=1e-2) + assert_allclose(out.affine, dst_affine, rtol=1e-2, atol=1e-2) + + inverted = xform.inverse(output_data)["img"] + self.assertEqual(inverted.applied_operations, []) # no further invert after inverting + expected_affine = to_affine_nd(len(out.affine) - 1, torch.eye(4)) + assert_allclose(inverted.affine, expected_affine, rtol=1e-2, atol=1e-2) + assert_allclose(inverted, img, rtol=1e-2, atol=1e-2) if __name__ == "__main__": diff --git a/tests/test_split_channel.py b/tests/test_split_channel.py index 75216227e4..4b41c334e8 100644 --- a/tests/test_split_channel.py +++ b/tests/test_split_channel.py @@ -30,6 +30,7 @@ class TestSplitChannel(unittest.TestCase): def test_shape(self, input_param, test_data, expected_shape): result = SplitChannel(**input_param)(test_data) for data in result: + self.assertEqual(type(data), type(test_data)) self.assertTupleEqual(data.shape, expected_shape) diff --git a/tests/test_splitdim.py b/tests/test_splitdim.py new file mode 100644 index 0000000000..d6ee4fc55e --- /dev/null +++ b/tests/test_splitdim.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. + +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.transforms.utility.array import SplitDim +from tests.utils import TEST_NDARRAYS + +TESTS = [] +for p in TEST_NDARRAYS: + for keepdim in (True, False): + TESTS.append(((2, 10, 8, 7), keepdim, p)) + + +class TestSplitDim(unittest.TestCase): + @parameterized.expand(TESTS) + def test_correct_shape(self, shape, keepdim, im_type): + arr = im_type(np.random.rand(*shape)) + for dim in range(arr.ndim): + out = SplitDim(dim, keepdim)(arr) + self.assertIsInstance(out, (list, tuple)) + self.assertEqual(type(out[0]), type(arr)) + self.assertEqual(len(out), arr.shape[dim]) + expected_ndim = arr.ndim if keepdim else arr.ndim - 1 + self.assertEqual(out[0].ndim, expected_ndim) + # assert is a shallow copy + arr[0, 0, 0, 0] *= 2 + self.assertEqual(arr.flatten()[0], out[0].flatten()[0]) + + def test_error(self): + """Should fail because splitting along singleton dimension""" + shape = (2, 1, 8, 7) + for p in TEST_NDARRAYS: + arr = p(np.random.rand(*shape)) + with self.assertRaises(RuntimeError): + _ = SplitDim(dim=1)(arr) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_splitdimd.py b/tests/test_splitdimd.py new file mode 100644 index 0000000000..1e39439b86 --- /dev/null +++ b/tests/test_splitdimd.py @@ -0,0 +1,82 @@ +# 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. + +import unittest +from copy import deepcopy + +import numpy as np +import torch +from parameterized import parameterized + +from monai.data.meta_tensor import MetaTensor +from monai.transforms import LoadImaged +from monai.transforms.utility.dictionary import SplitDimd +from tests.utils import TEST_NDARRAYS, assert_allclose, make_nifti_image, make_rand_affine + +TESTS = [] +for p in TEST_NDARRAYS: + for keepdim in (True, False): + for update_meta in (True, False): + TESTS.append((keepdim, p, update_meta)) + + +class TestSplitDimd(unittest.TestCase): + @classmethod + def setUpClass(cls): + arr = np.random.rand(2, 10, 8, 7) + affine = make_rand_affine() + data = {"i": make_nifti_image(arr, affine)} + + loader = LoadImaged("i") + cls.data: MetaTensor = loader(data) + + @parameterized.expand(TESTS) + def test_correct(self, keepdim, im_type, update_meta): + data = deepcopy(self.data) + data["i"] = im_type(data["i"]) + arr = data["i"] + for dim in range(arr.ndim): + out = SplitDimd("i", dim=dim, keepdim=keepdim, update_meta=update_meta)(data) + self.assertIsInstance(out, dict) + self.assertEqual(len(out.keys()), len(data.keys()) + arr.shape[dim]) + # if updating metadata, pick some random points and + # check same world coordinates between input and output + if update_meta: + for _ in range(10): + idx = [np.random.choice(i) for i in arr.shape] + split_im_idx = idx[dim] + split_idx = deepcopy(idx) + split_idx[dim] = 0 + split_im = out[f"i_{split_im_idx}"] + if isinstance(data, MetaTensor) and isinstance(split_im, MetaTensor): + # idx[1:] to remove channel and then add 1 for 4th element + real_world = data.affine @ torch.tensor(idx[1:] + [1]).double() + real_world2 = split_im.affine @ torch.tensor(split_idx[1:] + [1]).double() + assert_allclose(real_world, real_world2) + + out = out["i_0"] + expected_ndim = arr.ndim if keepdim else arr.ndim - 1 + self.assertEqual(out.ndim, expected_ndim) + # assert is a shallow copy + arr[0, 0, 0, 0] *= 2 + self.assertEqual(arr.flatten()[0], out.flatten()[0]) + + def test_error(self): + """Should fail because splitting along singleton dimension""" + shape = (2, 1, 8, 7) + for p in TEST_NDARRAYS: + arr = p(np.random.rand(*shape)) + with self.assertRaises(RuntimeError): + _ = SplitDimd("i", dim=1)({"i": arr}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_std_shift_intensity.py b/tests/test_std_shift_intensity.py index 55750161ec..a5549bf187 100644 --- a/tests/test_std_shift_intensity.py +++ b/tests/test_std_shift_intensity.py @@ -15,6 +15,7 @@ import torch from monai.transforms import ShiftIntensity, StdShiftIntensity +from monai.utils import dtype_numpy_to_torch from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D @@ -67,7 +68,7 @@ def test_dtype(self): factor = np.random.rand() std_shifter = StdShiftIntensity(factor=factor, dtype=trans_dtype) result = std_shifter(image) - np.testing.assert_equal(result.dtype, trans_dtype) + np.testing.assert_equal(result.dtype, dtype_numpy_to_torch(trans_dtype)) if __name__ == "__main__": diff --git a/tests/test_std_shift_intensityd.py b/tests/test_std_shift_intensityd.py index 595da5cbc2..b86f6bd5e6 100644 --- a/tests/test_std_shift_intensityd.py +++ b/tests/test_std_shift_intensityd.py @@ -14,6 +14,7 @@ import numpy as np from monai.transforms import ShiftIntensityd, StdShiftIntensityd +from monai.utils import dtype_numpy_to_torch from tests.utils import NumpyImageTestCase2D @@ -64,7 +65,7 @@ def test_dtype(self): factor = np.random.rand() std_shifter = StdShiftIntensityd(keys=[key], factor=factor, dtype=trans_dtype) result = std_shifter({key: image}) - np.testing.assert_equal(result[key].dtype, trans_dtype) + np.testing.assert_equal(result[key].dtype, dtype_numpy_to_torch(trans_dtype)) if __name__ == "__main__": diff --git a/tests/test_subpixel_upsample.py b/tests/test_subpixel_upsample.py index 0216f164c3..3e5370473c 100644 --- a/tests/test_subpixel_upsample.py +++ b/tests/test_subpixel_upsample.py @@ -18,6 +18,7 @@ from monai.networks import eval_mode from monai.networks.blocks import SubpixelUpsample from monai.networks.layers.factories import Conv +from tests.utils import SkipIfBeforePyTorchVersion, test_script_save TEST_CASE_SUBPIXEL = [] for inch in range(1, 5): @@ -73,6 +74,13 @@ def test_subpixel_shape(self, input_param, input_shape, expected_shape): result = net.forward(torch.randn(input_shape)) self.assertEqual(result.shape, expected_shape) + @SkipIfBeforePyTorchVersion((1, 8, 1)) + def test_script(self): + input_param, input_shape, _ = TEST_CASE_SUBPIXEL[0] + net = SubpixelUpsample(**input_param) + test_data = torch.randn(input_shape) + test_script_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_surface_dice.py b/tests/test_surface_dice.py new file mode 100644 index 0000000000..ccc6242e1e --- /dev/null +++ b/tests/test_surface_dice.py @@ -0,0 +1,292 @@ +# 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. + +import unittest + +import numpy as np +import torch +import torch.nn.functional as F + +from monai.metrics.surface_dice import SurfaceDiceMetric + + +class TestAllSurfaceDiceMetrics(unittest.TestCase): + def test_tolerance_euclidean_distance(self): + batch_size = 2 + n_class = 2 + predictions = torch.zeros((batch_size, 480, 640), dtype=torch.int64) + labels = torch.zeros((batch_size, 480, 640), dtype=torch.int64) + predictions[0, :, 50:] = 1 + labels[0, :, 60:] = 1 # 10 px shift + predictions_hot = F.one_hot(predictions, num_classes=n_class).permute(0, 3, 1, 2) + labels_hot = F.one_hot(labels, num_classes=n_class).permute(0, 3, 1, 2) + + sd0 = SurfaceDiceMetric(class_thresholds=[0, 0], include_background=True) + res0 = sd0(predictions_hot, labels_hot) + agg0 = sd0.aggregate() # aggregation: nanmean across image then nanmean across batch + sd0_nans = SurfaceDiceMetric(class_thresholds=[0, 0], include_background=True, get_not_nans=True) + res0_nans = sd0_nans(predictions_hot, labels_hot) + agg0_nans, not_nans = sd0_nans.aggregate() + + np.testing.assert_array_equal(res0, res0_nans) + np.testing.assert_array_equal(agg0, agg0_nans) + + res1 = SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, labels_hot) + res9 = SurfaceDiceMetric(class_thresholds=[9, 9], include_background=True)(predictions_hot, labels_hot) + res10 = SurfaceDiceMetric(class_thresholds=[10, 10], include_background=True)(predictions_hot, labels_hot) + res11 = SurfaceDiceMetric(class_thresholds=[11, 11], include_background=True)(predictions_hot, labels_hot) + + for res in [res0, res9, res10, res11]: + assert res.shape == torch.Size([2, 2]) + + assert res0[0, 0] < res1[0, 0] < res9[0, 0] < res10[0, 0] + assert res0[0, 1] < res1[0, 1] < res9[0, 1] < res10[0, 1] + np.testing.assert_array_equal(res10, res11) + + expected_res0 = np.zeros((batch_size, n_class)) + expected_res0[0, 1] = 1 - (478 + 480 + 9 * 2) / (480 * 4 + 588 * 2 + 578 * 2) + expected_res0[0, 0] = 1 - (478 + 480 + 9 * 2) / (480 * 4 + 48 * 2 + 58 * 2) + expected_res0[1, 0] = 1 + expected_res0[1, 1] = np.nan + for b, c in np.ndindex(batch_size, n_class): + np.testing.assert_allclose(expected_res0[b, c], res0[b, c]) + np.testing.assert_array_equal(agg0, np.nanmean(np.nanmean(expected_res0, axis=1), axis=0)) + np.testing.assert_equal(not_nans, torch.tensor(2)) + + def test_tolerance_all_distances(self): + batch_size = 1 + n_class = 2 + predictions = torch.zeros((batch_size, 10, 10), dtype=torch.int64) + labels = torch.zeros((batch_size, 10, 10), dtype=torch.int64) + predictions[0, 1:4, 1] = 1 + """ + [[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [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, 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, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]] + """ + labels[0, 5:8, 6] = 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, 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, 0, 0, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 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]]] + """ + predictions_hot = F.one_hot(predictions, num_classes=n_class).permute(0, 3, 1, 2) + labels_hot = F.one_hot(labels, num_classes=n_class).permute(0, 3, 1, 2) + + # Euclidean distance: + # background: + # 36 boundary pixels have 0 distances; non-zero distances: + # distances gt_pred: [3, np.sqrt(9+4), 2, 3, 2, 2, 2, 1] + # distances pred_gt: [1, 2, 2, 1] + # class 1: + # distances gt_pred: [sqrt(25+4), sqrt(25+9), sqrt(25+16)] = [5.38516481, 5.83095189, 6.40312424] + # distances pred_gt: [sqrt(25+16), sqrt(25+9), sqrt(25+4)] = [6.40312424, 5.83095189, 5.38516481] + + res = SurfaceDiceMetric(class_thresholds=[0, 0], include_background=True)(predictions_hot, labels_hot) + expected_res = [[1 - (8 + 4) / (36 * 2 + 8 + 4), 0]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[2.8, 5.5], include_background=True)(predictions_hot, labels_hot) + expected_res = [[1 - 3 / (36 * 2 + 8 + 4), 1 - (2 + 2) / (3 + 3)]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[3, 6], include_background=True)(predictions_hot, labels_hot) + expected_res = [[1 - 1 / (36 * 2 + 8 + 4), 1 - 2 / (3 + 3)]] + np.testing.assert_array_almost_equal(res, expected_res) + + # Chessboard distance: + # background: + # 36 boundary pixels have 0 distances; non-zero distances: + # distances gt_pred: [max(3,0), max(3,2), max(2,0), max(3,3), max(2,0), max(0,2), max(2,0), max(0,1)] = + # [3, 3, 2, 3, 2, 2, 2, 1] + # distances pred_gt: [max(0,1), max(2,0), max(2,0), max(1,0)] = [1, 2, 2, 1] + # class 1: + # distances gt_pred: [max(5,2), max(5,3), max(5,4)] = [5, 5, 5] + # distances pred_gt: [max(5,4), max(5,3), max(5,2)] = [5, 5, 5] + + res = SurfaceDiceMetric(class_thresholds=[0, 0], include_background=True, distance_metric="chessboard")( + predictions_hot, labels_hot + ) + expected_res = [[1 - (8 + 4) / (36 * 2 + 8 + 4), 0]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[1, 4.999], include_background=True, distance_metric="chessboard")( + predictions_hot, labels_hot + ) + expected_res = [[1 - (7 + 2) / (36 * 2 + 8 + 4), 0]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[2, 5], include_background=True, distance_metric="chessboard")( + predictions_hot, labels_hot + ) + expected_res = [[1 - 3 / (36 * 2 + 8 + 4), 1]] + np.testing.assert_array_almost_equal(res, expected_res) + + # Taxicab distance (= Manhattan distance): + # background: + # 36 boundary pixels have 0 distances; non-zero distances: + # distances gt_pred: [3+0, 4+0, 2+0, 0+3, 2+0, 0+2, 2+0, 0+1] = [3, 4, 2, 3, 2, 2, 2, 1] + # distances pred_gt: [0+1, 2+0, 2+0, 1+0] = [1, 2, 2, 1] + # class 1: + # distances gt_pred: [5+2, 5+3, 5+4] = [7, 8, 9] + # distances pred_gt: [5+4, 5+3, 5+2] = [9, 8, 7] + + res = SurfaceDiceMetric(class_thresholds=[0, 0], include_background=True, distance_metric="taxicab")( + predictions_hot, labels_hot + ) + expected_res = [[1 - (8 + 4) / (36 * 2 + 8 + 4), 0]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[1, 7], include_background=True, distance_metric="taxicab")( + predictions_hot, labels_hot + ) + expected_res = [[1 - (7 + 2) / (36 * 2 + 8 + 4), 1 - (2 + 2) / (3 + 3)]] + np.testing.assert_array_almost_equal(res, expected_res) + + res = SurfaceDiceMetric(class_thresholds=[3, 9], include_background=True, distance_metric="taxicab")( + predictions_hot, labels_hot + ) + expected_res = [[1 - 1 / (36 * 2 + 8 + 4), 1]] + np.testing.assert_array_almost_equal(res, expected_res) + + def test_asserts(self): + batch_size = 1 + n_class = 2 + predictions = torch.zeros((batch_size, 80, 80), dtype=torch.int64) + labels = torch.zeros((batch_size, 80, 80), dtype=torch.int64) + predictions[0, 10:20, 10:20] = 1 + labels[0, 20:30, 20:30] = 1 + predictions_hot = F.one_hot(predictions, num_classes=n_class).permute(0, 3, 1, 2) + labels_hot = F.one_hot(labels, num_classes=n_class).permute(0, 3, 1, 2) + + # no torch tensor + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot.numpy(), labels_hot) + self.assertEqual( + "y_pred or y must be a list/tuple of `channel-first` Tensors or a `batch-first` Tensor.", + str(context.exception), + ) + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, labels_hot.numpy()) + self.assertEqual("y_pred and y must be PyTorch Tensor.", str(context.exception)) + + # wrong dimensions + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions, labels_hot) + self.assertEqual("y_pred and y should have four dimensions: [B,C,H,W].", str(context.exception)) + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, labels) + self.assertEqual("y_pred and y should have four dimensions: [B,C,H,W].", str(context.exception)) + + # mismatch of shape of input tensors + input_bad_shape = torch.clone(predictions_hot) + input_bad_shape = input_bad_shape[:, :, :, :50] + + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, input_bad_shape) + self.assertEqual( + "y_pred and y should have same shape, but instead, shapes are torch.Size([1, 2, 80, 80]) (y_pred) and " + "torch.Size([1, 2, 80, 50]) (y).", + str(context.exception), + ) + + # input tensors not one-hot encoded + predictions_no_hot = torch.clone(predictions_hot) + predictions_no_hot[0, :, 0, 0] = torch.tensor([2, 0]) + + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_no_hot, predictions_hot) + self.assertEqual("y_pred and y should be one-hot encoded.", str(context.exception)) + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, predictions_no_hot) + self.assertEqual("y_pred and y should be one-hot encoded.", str(context.exception)) + + predictions_no_hot = predictions_no_hot.float() + predictions_no_hot[0, :, 0, 0] = torch.tensor([0.5, 0]) + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_no_hot, predictions_hot) + self.assertEqual("y_pred and y should be binarized tensors (e.g. torch.int64).", str(context.exception)) + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1], include_background=True)(predictions_hot, predictions_no_hot) + self.assertEqual("y_pred and y should be binarized tensors (e.g. torch.int64).", str(context.exception)) + + # wrong number of class thresholds + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[1, 1, 1], include_background=True)(predictions_hot, labels_hot) + self.assertEqual("number of classes (2) does not match number of class thresholds (3).", str(context.exception)) + + # inf and nan values in class thresholds + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[np.inf, 1], include_background=True)(predictions_hot, labels_hot) + self.assertEqual("All class thresholds need to be finite.", str(context.exception)) + + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[np.nan, 1], include_background=True)(predictions_hot, labels_hot) + self.assertEqual("All class thresholds need to be finite.", str(context.exception)) + + # negative values in class thresholds: + with self.assertRaises(ValueError) as context: + SurfaceDiceMetric(class_thresholds=[-0.22, 1], include_background=True)(predictions_hot, labels_hot) + self.assertEqual("All class thresholds need to be >= 0.", str(context.exception)) + + def test_not_predicted_not_present(self): + # class is present in labels, but not in prediction -> nsd of 0 should be yielded for that class; class is + # neither present on labels, nor prediction -> nan should be yielded + batch_size = 1 + n_class = 4 + predictions = torch.zeros((batch_size, 80, 80), dtype=torch.int64) + labels = torch.zeros((batch_size, 80, 80), dtype=torch.int64) + predictions[0, 10:20, 10:20] = 1 + labels[0, 10:20, 10:20] = 2 + predictions_hot = F.one_hot(predictions, num_classes=n_class).permute(0, 3, 1, 2) + labels_hot = F.one_hot(labels, num_classes=n_class).permute(0, 3, 1, 2) + + # with and without background class + sur_metric_bgr = SurfaceDiceMetric(class_thresholds=[1, 1, 1, 1], include_background=True) + sur_metric = SurfaceDiceMetric(class_thresholds=[1, 1, 1], include_background=False) + + # test per-class results + res_bgr_classes = sur_metric_bgr(predictions_hot, labels_hot) + np.testing.assert_array_equal(res_bgr_classes, [[1, 0, 0, np.nan]]) + res_classes = sur_metric(predictions_hot, labels_hot) + np.testing.assert_array_equal(res_classes, [[0, 0, np.nan]]) + + # test aggregation + res_bgr = sur_metric_bgr.aggregate(reduction="mean") + np.testing.assert_equal(res_bgr, torch.tensor([1 / 3], dtype=torch.float64)) + res = sur_metric.aggregate() + np.testing.assert_equal(res, torch.tensor([0], dtype=torch.float64)) + + predictions_empty = torch.zeros((2, 3, 1, 1)) + sur_metric_nans = SurfaceDiceMetric(class_thresholds=[1, 1, 1], include_background=True, get_not_nans=True) + res_classes = sur_metric_nans(predictions_empty, predictions_empty) + res, not_nans = sur_metric_nans.aggregate() + np.testing.assert_array_equal(res_classes, [[np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan]]) + np.testing.assert_equal(res, torch.tensor([0], dtype=torch.float64)) + np.testing.assert_equal(not_nans, torch.tensor([0], dtype=torch.float64)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_surface_distance.py b/tests/test_surface_distance.py index edfe9e8663..4cd70b43aa 100644 --- a/tests/test_surface_distance.py +++ b/tests/test_surface_distance.py @@ -134,7 +134,7 @@ def test_nans(self, input_data): batch_seg_1 = [seg_1.unsqueeze(0)] batch_seg_2 = [seg_2.unsqueeze(0)] sur_metric(batch_seg_1, batch_seg_2) - result, not_nans = sur_metric.aggregate() + result, not_nans = sur_metric.aggregate(reduction="mean") np.testing.assert_allclose(0, result, rtol=1e-5) np.testing.assert_allclose(0, not_nans, rtol=1e-5) diff --git a/tests/test_swin_unetr.py b/tests/test_swin_unetr.py new file mode 100644 index 0000000000..0d48e99c44 --- /dev/null +++ b/tests/test_swin_unetr.py @@ -0,0 +1,89 @@ +# 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. + +import unittest +from unittest import skipUnless + +import torch +from parameterized import parameterized + +from monai.networks import eval_mode +from monai.networks.nets.swin_unetr import SwinUNETR +from monai.utils import optional_import + +einops, has_einops = optional_import("einops") + +TEST_CASE_SWIN_UNETR = [] +for attn_drop_rate in [0.4]: + for in_channels in [1]: + for depth in [[2, 1, 1, 1], [1, 2, 1, 1]]: + for out_channels in [2]: + for img_size in [64]: + for feature_size in [12]: + for norm_name in ["instance"]: + for nd in (2, 3): + test_case = [ + { + "in_channels": in_channels, + "out_channels": out_channels, + "img_size": (img_size,) * nd, + "feature_size": feature_size, + "depths": depth, + "norm_name": norm_name, + "attn_drop_rate": attn_drop_rate, + }, + (2, in_channels, *([img_size] * nd)), + (2, out_channels, *([img_size] * nd)), + ] + if nd == 2: + test_case[0]["spatial_dims"] = 2 # type: ignore + TEST_CASE_SWIN_UNETR.append(test_case) + + +class TestSWINUNETR(unittest.TestCase): + @parameterized.expand(TEST_CASE_SWIN_UNETR) + @skipUnless(has_einops, "Requires einops") + def test_shape(self, input_param, input_shape, expected_shape): + net = SwinUNETR(**input_param) + with eval_mode(net): + result = net(torch.randn(input_shape)) + self.assertEqual(result.shape, expected_shape) + + def test_ill_arg(self): + with self.assertRaises(ValueError): + SwinUNETR( + in_channels=1, + out_channels=3, + img_size=(128, 128, 128), + feature_size=24, + norm_name="instance", + attn_drop_rate=4, + ) + + with self.assertRaises(ValueError): + SwinUNETR(in_channels=1, out_channels=2, img_size=(96, 96), feature_size=48, norm_name="instance") + + with self.assertRaises(ValueError): + SwinUNETR(in_channels=1, out_channels=4, img_size=(96, 96, 96), feature_size=50, norm_name="instance") + + with self.assertRaises(ValueError): + SwinUNETR( + in_channels=1, + out_channels=3, + img_size=(85, 85, 85), + feature_size=24, + norm_name="instance", + drop_rate=0.4, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_testtimeaugmentation.py b/tests/test_testtimeaugmentation.py index 21186adc3c..d5aa1af688 100644 --- a/tests/test_testtimeaugmentation.py +++ b/tests/test_testtimeaugmentation.py @@ -32,7 +32,7 @@ RandScaleIntensityd, ) from monai.transforms.croppad.dictionary import SpatialPadd -from monai.transforms.spatial.dictionary import RandFlipd, Spacingd +from monai.transforms.spatial.dictionary import RandFlipd from monai.utils import optional_import, set_determinism from monai.utils.enums import PostFix from tests.utils import TEST_NDARRAYS @@ -58,9 +58,7 @@ def get_data(num_examples, input_size, data_type=np.asarray, include_label=True) data = [] for i in range(num_examples): im, label = custom_create_test_image_2d() - d = {} - d["image"] = data_type(im[:, i:]) - d[PostFix.meta("image")] = {"affine": np.eye(4)} + d = {"image": data_type(im[:, i:])} if include_label: d["label"] = data_type(label[:, i:]) d[PostFix.meta("label")] = {"affine": np.eye(4)} @@ -176,12 +174,6 @@ def test_image_no_label(self): tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, orig_key="image") tta(self.get_data(1, (20, 20), include_label=False)) - @unittest.skipUnless(has_nib, "Requires nibabel") - def test_requires_meta_dict(self): - transforms = Compose([AddChanneld("image"), RandFlipd("image"), Spacingd("image", pixdim=1.1)]) - tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, orig_key="image") - tta(self.get_data(1, (20, 20), include_label=False)) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_thread_buffer.py b/tests/test_thread_buffer.py index 04511220f8..013c20f4ce 100644 --- a/tests/test_thread_buffer.py +++ b/tests/test_thread_buffer.py @@ -94,6 +94,16 @@ def test_dataloader_repeats(self): self.assertTrue(previous_batch is d, "Batch object was not repeated") previous_batch = None + def test_thread_workers(self): + dataset = Dataset(data=self.datalist, transform=self.transform) + dataloader = ThreadDataLoader(dataset=dataset, batch_size=2, num_workers=2, use_thread_workers=True) + + for d in dataloader: + self.assertEqual(d["image"][0], "spleen_19.nii.gz") + self.assertEqual(d["image"][1], "spleen_31.nii.gz") + self.assertEqual(d["label"][0], "spleen_label_19.nii.gz") + self.assertEqual(d["label"][1], "spleen_label_31.nii.gz") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_threshold_intensity.py b/tests/test_threshold_intensity.py index 01321f1b0b..3c0a2033ee 100644 --- a/tests/test_threshold_intensity.py +++ b/tests/test_threshold_intensity.py @@ -29,7 +29,7 @@ class TestThresholdIntensity(unittest.TestCase): def test_value(self, in_type, input_param, expected_value): test_data = in_type(np.arange(10)) result = ThresholdIntensity(**input_param)(test_data) - assert_allclose(result, in_type(expected_value)) + assert_allclose(result, in_type(expected_value), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_threshold_intensityd.py b/tests/test_threshold_intensityd.py index e0610ebb5b..8aade12322 100644 --- a/tests/test_threshold_intensityd.py +++ b/tests/test_threshold_intensityd.py @@ -47,9 +47,9 @@ class TestThresholdIntensityd(unittest.TestCase): def test_value(self, in_type, input_param, expected_value): test_data = {"image": in_type(np.arange(10)), "label": in_type(np.arange(10)), "extra": in_type(np.arange(10))} result = ThresholdIntensityd(**input_param)(test_data) - assert_allclose(result["image"], in_type(expected_value)) - assert_allclose(result["label"], in_type(expected_value)) - assert_allclose(result["extra"], in_type(expected_value)) + assert_allclose(result["image"], in_type(expected_value), type_test="tensor") + assert_allclose(result["label"], in_type(expected_value), type_test="tensor") + assert_allclose(result["extra"], in_type(expected_value), type_test="tensor") if __name__ == "__main__": diff --git a/tests/test_to_from_meta_tensord.py b/tests/test_to_from_meta_tensord.py new file mode 100644 index 0000000000..6f46055d6a --- /dev/null +++ b/tests/test_to_from_meta_tensord.py @@ -0,0 +1,176 @@ +# 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. + +import random +import string +import unittest +from copy import deepcopy +from typing import Optional, Union + +import torch +from parameterized import parameterized + +from monai.data.meta_tensor import MetaTensor +from monai.transforms import FromMetaTensord, ToMetaTensord +from monai.utils.enums import PostFix +from monai.utils.module import get_torch_version_tuple +from tests.utils import TEST_DEVICES, assert_allclose + +PT_VER_MAJ, PT_VER_MIN = get_torch_version_tuple() + +DTYPES = [[torch.float32], [torch.float64], [torch.float16], [torch.int64], [torch.int32]] +TESTS = [] +for _device in TEST_DEVICES: + for _dtype in DTYPES: + TESTS.append((*_device, *_dtype)) + + +def rand_string(min_len=5, max_len=10): + str_size = random.randint(min_len, max_len) + chars = string.ascii_letters + string.punctuation + return "".join(random.choice(chars) for _ in range(str_size)) + + +class TestToFromMetaTensord(unittest.TestCase): + @staticmethod + def get_im(shape=None, dtype=None, device=None): + if shape is None: + shape = shape = (1, 10, 8) + affine = torch.randint(0, 10, (4, 4)) + meta = {"fname": rand_string()} + t = torch.rand(shape) + if dtype is not None: + t = t.to(dtype) + if device is not None: + t = t.to(device) + m = MetaTensor(t.clone(), affine, meta) + return m + + def check_ids(self, a, b, should_match): + comp = self.assertEqual if should_match else self.assertNotEqual + comp(id(a), id(b)) + + def check( + self, + out: torch.Tensor, + orig: torch.Tensor, + *, + shape: bool = True, + vals: bool = True, + ids: bool = True, + device: Optional[Union[str, torch.device]] = None, + meta: bool = True, + check_ids: bool = True, + **kwargs, + ): + if device is None: + device = orig.device + + # check the image + self.assertIsInstance(out, type(orig)) + if shape: + assert_allclose(torch.as_tensor(out.shape), torch.as_tensor(orig.shape)) + if vals: + assert_allclose(out, orig, **kwargs) + if check_ids: + self.check_ids(out, orig, ids) + self.assertTrue(str(device) in str(out.device)) + + # check meta and affine are equal and affine is on correct device + if isinstance(orig, MetaTensor) and isinstance(out, MetaTensor) and meta: + orig_meta_no_affine = deepcopy(orig.meta) + del orig_meta_no_affine["affine"] + out_meta_no_affine = deepcopy(out.meta) + del out_meta_no_affine["affine"] + self.assertEqual(orig_meta_no_affine, out_meta_no_affine) + assert_allclose(out.affine, orig.affine) + self.assertTrue(str(device) in str(out.affine.device)) + if check_ids: + self.check_ids(out.affine, orig.affine, ids) + self.check_ids(out.meta, orig.meta, ids) + + @parameterized.expand(TESTS) + def test_from_to_meta_tensord(self, device, dtype): + m1 = self.get_im(device=device, dtype=dtype) + m2 = self.get_im(device=device, dtype=dtype) + m3 = self.get_im(device=device, dtype=dtype) + d_metas = {"m1": m1, "m2": m2, "m3": m3} + m1_meta = {k: v for k, v in m1.meta.items() if k != "affine"} + m1_aff = m1.affine + + # FROM -> forward + t_from_meta = FromMetaTensord(["m1", "m2"]) + d_dict = t_from_meta(d_metas) + + self.assertEqual( + sorted(d_dict.keys()), + [ + "m1", + PostFix.meta("m1"), + PostFix.transforms("m1"), + "m2", + PostFix.meta("m2"), + PostFix.transforms("m2"), + "m3", + ], + ) + self.check(d_dict["m3"], m3, ids=True) # unchanged + self.check(d_dict["m1"], m1.as_tensor(), ids=False) + meta_out = {k: v for k, v in d_dict["m1_meta_dict"].items() if k != "affine"} + aff_out = d_dict["m1_meta_dict"]["affine"] + self.check(aff_out, m1_aff, ids=False) + self.assertEqual(meta_out, m1_meta) + + # FROM -> inverse + d_meta_dict_meta = t_from_meta.inverse(d_dict) + self.assertEqual(sorted(d_meta_dict_meta.keys()), ["m1", "m2", "m3"]) + self.check(d_meta_dict_meta["m3"], m3, ids=False) # unchanged (except deep copy in inverse) + self.check(d_meta_dict_meta["m1"], m1, ids=False) + meta_out = {k: v for k, v in d_meta_dict_meta["m1"].meta.items() if k != "affine"} + aff_out = d_meta_dict_meta["m1"].affine + self.check(aff_out, m1_aff, ids=False) + self.assertEqual(meta_out, m1_meta) + + # TO -> Forward + t_to_meta = ToMetaTensord(["m1", "m2"]) + d_dict_meta = t_to_meta(d_dict) + self.assertEqual(sorted(d_dict_meta.keys()), ["m1", "m2", "m3"]) + self.check(d_dict_meta["m3"], m3, ids=True) # unchanged (except deep copy in inverse) + self.check(d_dict_meta["m1"], m1, ids=False) + meta_out = {k: v for k, v in d_dict_meta["m1"].meta.items() if k != "affine"} + aff_out = d_dict_meta["m1"].meta["affine"] + self.check(aff_out, m1_aff, ids=False) + self.assertEqual(meta_out, m1_meta) + + # TO -> Inverse + d_dict_meta_dict = t_to_meta.inverse(d_dict_meta) + self.assertEqual( + sorted(d_dict_meta_dict.keys()), + [ + "m1", + PostFix.meta("m1"), + PostFix.transforms("m1"), + "m2", + PostFix.meta("m2"), + PostFix.transforms("m2"), + "m3", + ], + ) + self.check(d_dict_meta_dict["m3"], m3.as_tensor(), ids=False) # unchanged (except deep copy in inverse) + self.check(d_dict_meta_dict["m1"], m1.as_tensor(), ids=False) + meta_out = {k: v for k, v in d_dict_meta_dict["m1_meta_dict"].items() if k != "affine"} + aff_out = d_dict_meta_dict["m1_meta_dict"]["affine"] + self.check(aff_out, m1_aff, ids=False) + self.assertEqual(meta_out, m1_meta) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_to_tensor.py b/tests/test_to_tensor.py index bfc61cdb19..fe6c4fb0d4 100644 --- a/tests/test_to_tensor.py +++ b/tests/test_to_tensor.py @@ -21,13 +21,11 @@ im = [[1, 2], [3, 4]] -TESTS = [] -TESTS.append((im, (2, 2))) +TESTS = [(im, (2, 2))] for p in TEST_NDARRAYS: TESTS.append((p(im), (2, 2))) -TESTS_SINGLE = [] -TESTS_SINGLE.append([5]) +TESTS_SINGLE = [[5]] for p in TEST_NDARRAYS: TESTS_SINGLE.append([p(5)]) diff --git a/tests/test_torchscript_utils.py b/tests/test_torchscript_utils.py index d6bea09ed6..cdf2f19eb3 100644 --- a/tests/test_torchscript_utils.py +++ b/tests/test_torchscript_utils.py @@ -18,7 +18,6 @@ from monai.config import get_config_values from monai.data import load_net_with_metadata, save_net_with_metadata from monai.utils import JITMetadataKeys -from monai.utils.module import pytorch_after class TestModule(torch.nn.Module): @@ -102,10 +101,7 @@ def test_save_load_more_extra_files(self): _, _, loaded_extra_files = load_net_with_metadata(f"{tempdir}/test.ts", more_extra_files=("test.txt",)) - if pytorch_after(1, 7): - self.assertEqual(more_extra_files["test.txt"], loaded_extra_files["test.txt"]) - else: - self.assertEqual(more_extra_files["test.txt"].decode(), loaded_extra_files["test.txt"]) + self.assertEqual(more_extra_files["test.txt"], loaded_extra_files["test.txt"]) if __name__ == "__main__": diff --git a/tests/test_torchvision.py b/tests/test_torchvision.py index e0844eb4b9..68b9413e65 100644 --- a/tests/test_torchvision.py +++ b/tests/test_torchvision.py @@ -15,7 +15,7 @@ from monai.transforms import TorchVision from monai.utils import set_determinism -from tests.utils import TEST_NDARRAYS, SkipIfBeforePyTorchVersion, assert_allclose +from tests.utils import TEST_NDARRAYS, assert_allclose TESTS = [] for p in TEST_NDARRAYS: @@ -52,7 +52,6 @@ ) -@SkipIfBeforePyTorchVersion((1, 7)) class TestTorchVision(unittest.TestCase): @parameterized.expand(TESTS) def test_value(self, input_param, input_data, expected_value): diff --git a/tests/test_torchvision_fully_conv_model.py b/tests/test_torchvision_fully_conv_model.py deleted file mode 100644 index 34a61ce9fa..0000000000 --- a/tests/test_torchvision_fully_conv_model.py +++ /dev/null @@ -1,81 +0,0 @@ -# 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. - -import unittest -from unittest import skipUnless - -import torch -from parameterized import parameterized - -from monai.networks import eval_mode -from monai.networks.nets import TorchVisionFullyConvModel -from monai.utils import optional_import - -_, has_tv = optional_import("torchvision") - -device = "cuda" if torch.cuda.is_available() else "cpu" - -TEST_CASE_0 = [{"model_name": "resnet18", "num_classes": 1, "pretrained": False}, (2, 3, 224, 224), (2, 1, 1, 1)] - -TEST_CASE_1 = [{"model_name": "resnet18", "num_classes": 1, "pretrained": False}, (2, 3, 256, 256), (2, 1, 2, 2)] - -TEST_CASE_2 = [{"model_name": "resnet101", "num_classes": 5, "pretrained": False}, (2, 3, 256, 256), (2, 5, 2, 2)] - -TEST_CASE_3 = [ - {"model_name": "resnet101", "num_classes": 5, "pool_size": 6, "pretrained": False}, - (2, 3, 224, 224), - (2, 5, 2, 2), -] - -TEST_CASE_PRETRAINED_0 = [ - {"model_name": "resnet18", "num_classes": 1, "pretrained": True}, - (2, 3, 224, 224), - (2, 1, 1, 1), - -0.010419349186122417, -] - -TEST_CASE_PRETRAINED_1 = [ - {"model_name": "resnet18", "num_classes": 1, "pretrained": True}, - (2, 3, 256, 256), - (2, 1, 2, 2), - -0.010419349186122417, -] - -TEST_CASE_PRETRAINED_2 = [ - {"model_name": "resnet18", "num_classes": 5, "pretrained": True}, - (2, 3, 256, 256), - (2, 5, 2, 2), - -0.010419349186122417, -] - - -class TestTorchVisionFullyConvModel(unittest.TestCase): - @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) - @skipUnless(has_tv, "Requires TorchVision.") - def test_without_pretrained(self, input_param, input_shape, expected_shape): - net = TorchVisionFullyConvModel(**input_param).to(device) - with eval_mode(net): - result = net.forward(torch.randn(input_shape).to(device)) - self.assertEqual(result.shape, expected_shape) - - @parameterized.expand([TEST_CASE_PRETRAINED_0, TEST_CASE_PRETRAINED_1, TEST_CASE_PRETRAINED_2]) - @skipUnless(has_tv, "Requires TorchVision.") - def test_with_pretrained(self, input_param, input_shape, expected_shape, expected_value): - net = TorchVisionFullyConvModel(**input_param).to(device) - with eval_mode(net): - result = net.forward(torch.randn(input_shape).to(device)) - value = next(net.parameters())[0, 0, 0, 0].item() - self.assertEqual(value, expected_value) - self.assertEqual(result.shape, expected_shape) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_torchvisiond.py b/tests/test_torchvisiond.py index 4c62c6e41a..def26fa26b 100644 --- a/tests/test_torchvisiond.py +++ b/tests/test_torchvisiond.py @@ -16,7 +16,6 @@ from monai.transforms import TorchVisiond from monai.utils import set_determinism -from tests.utils import SkipIfBeforePyTorchVersion TEST_CASE_1 = [ {"keys": "img", "name": "ColorJitter"}, @@ -49,7 +48,6 @@ ] -@SkipIfBeforePyTorchVersion((1, 7)) class TestTorchVisiond(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) def test_value(self, input_param, input_data, expected_value): diff --git a/tests/test_transchex.py b/tests/test_transchex.py index 462ce64fd6..713bc35f56 100644 --- a/tests/test_transchex.py +++ b/tests/test_transchex.py @@ -38,7 +38,7 @@ "num_classes": num_classes, "drop_out": drop_out, }, - (2, num_classes), # type: ignore + (2, num_classes), ] TEST_CASE_TRANSCHEX.append(test_case) diff --git a/tests/test_tversky_loss.py b/tests/test_tversky_loss.py index 2bb2409360..22f57cc8c6 100644 --- a/tests/test_tversky_loss.py +++ b/tests/test_tversky_loss.py @@ -16,7 +16,7 @@ from parameterized import parameterized from monai.losses import TverskyLoss -from tests.utils import SkipIfBeforePyTorchVersion, test_script_save +from tests.utils import test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -175,7 +175,6 @@ def test_input_warnings(self): loss = TverskyLoss(to_onehot_y=True) loss.forward(chn_input, chn_target) - @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = TverskyLoss() test_input = torch.ones(2, 1, 8, 8) diff --git a/tests/test_unet.py b/tests/test_unet.py index 5f126fed97..a90e32230b 100644 --- a/tests/test_unet.py +++ b/tests/test_unet.py @@ -96,6 +96,7 @@ "strides": (2, 2), "num_res_units": 1, "act": (Act.LEAKYRELU, {"negative_slope": 0.2}), + "adn_ordering": "NA", }, (16, 4, 32, 64, 48), (16, 3, 32, 64, 48), diff --git a/tests/test_unified_focal_loss.py b/tests/test_unified_focal_loss.py new file mode 100644 index 0000000000..1a7bb91059 --- /dev/null +++ b/tests/test_unified_focal_loss.py @@ -0,0 +1,63 @@ +# 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. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.losses import AsymmetricUnifiedFocalLoss + +TEST_CASES = [ + [ # 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, + ], + [ # 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, + ], +] + + +class TestAsymmetricUnifiedFocalLoss(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_result(self, input_data, expected_val): + loss = AsymmetricUnifiedFocalLoss() + result = loss(**input_data) + np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4) + + def test_ill_shape(self): + loss = AsymmetricUnifiedFocalLoss() + with self.assertRaisesRegex(ValueError, ""): + loss(torch.ones((2, 2, 2)), torch.ones((2, 2, 2, 2))) + + def test_with_cuda(self): + loss = AsymmetricUnifiedFocalLoss() + i = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]) + j = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]) + if torch.cuda.is_available(): + i = i.cuda() + j = j.cuda() + output = loss(i, j) + print(output) + np.testing.assert_allclose(output.detach().cpu().numpy(), 0.0, atol=1e-4, rtol=1e-4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_upsample_block.py b/tests/test_upsample_block.py index aa4141aabc..6e2e548042 100644 --- a/tests/test_upsample_block.py +++ b/tests/test_upsample_block.py @@ -76,8 +76,8 @@ test_case = [ {"dimensions": 3, "in_channels": 3, "out_channels": 5, "mode": t, "scale_factor": s, "align_corners": True}, (16, 3, 4, 5, 6), + expected_shape, ] - test_case.append(expected_shape) TEST_CASES_EQ.append(test_case) diff --git a/tests/test_utils_pytorch_numpy_unification.py b/tests/test_utils_pytorch_numpy_unification.py index b13378debe..df4db8a27c 100644 --- a/tests/test_utils_pytorch_numpy_unification.py +++ b/tests/test_utils_pytorch_numpy_unification.py @@ -12,12 +12,11 @@ import unittest import numpy as np -import torch from parameterized import parameterized from monai.transforms.utils_pytorch_numpy_unification import mode, percentile from monai.utils import set_determinism -from tests.utils import TEST_NDARRAYS, SkipIfBeforePyTorchVersion, assert_allclose +from tests.utils import TEST_NDARRAYS, assert_allclose TEST_MODE = [] for p in TEST_NDARRAYS: @@ -37,10 +36,7 @@ def test_percentile(self): for p in TEST_NDARRAYS: arr = p(np.arange(100 * 101).reshape(1, 100, 101).astype(np.float32)) results.append(percentile(arr, q)) - # pre torch 1.7, no `quantile`. Our own method doesn't interpolate, - # so we can only be accurate to 0.5 - atol = 0.5 if not hasattr(torch, "quantile") else 1e-4 - assert_allclose(results[0], results[-1], type_test=False, atol=atol) + assert_allclose(results[0], results[-1], type_test=False, atol=1e-4, rtol=1e-4) def test_fails(self): for p in TEST_NDARRAYS: @@ -49,17 +45,13 @@ def test_fails(self): with self.assertRaises(ValueError): percentile(arr, q) - @SkipIfBeforePyTorchVersion((1, 7)) def test_dim(self): q = np.random.randint(0, 100, size=50) results = [] for p in TEST_NDARRAYS: arr = p(np.arange(6).reshape(1, 2, 3).astype(np.float32)) results.append(percentile(arr, q, dim=1)) - # pre torch 1.7, no `quantile`. Our own method doesn't interpolate, - # so we can only be accurate to 0.5 - atol = 0.5 if not hasattr(torch, "quantile") else 1e-4 - assert_allclose(results[0], results[-1], type_test=False, atol=atol) + assert_allclose(results[0], results[-1], type_test=False, atol=1e-4) @parameterized.expand(TEST_MODE) def test_mode(self, array, expected, to_long): diff --git a/tests/test_varautoencoder.py b/tests/test_varautoencoder.py index 95fea8afcb..04fc07f53f 100644 --- a/tests/test_varautoencoder.py +++ b/tests/test_varautoencoder.py @@ -91,7 +91,7 @@ def test_script(self): spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=2, channels=(4, 8), strides=(2, 2) ) test_data = torch.randn(2, 1, 32, 32) - test_script_save(net, test_data) + test_script_save(net, test_data, rtol=1e-3, atol=1e-3) if __name__ == "__main__": diff --git a/tests/test_vis_gradbased.py b/tests/test_vis_gradbased.py new file mode 100644 index 0000000000..7655ca661e --- /dev/null +++ b/tests/test_vis_gradbased.py @@ -0,0 +1,50 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks.nets import DenseNet, DenseNet121, SEResNet50 +from monai.visualize import GuidedBackpropGrad, GuidedBackpropSmoothGrad, SmoothGrad, VanillaGrad + +DENSENET2D = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) +DENSENET3D = DenseNet(spatial_dims=3, in_channels=1, out_channels=3, init_features=2, growth_rate=2, block_config=(6,)) +SENET2D = SEResNet50(spatial_dims=2, in_channels=3, num_classes=4) +SENET3D = SEResNet50(spatial_dims=3, in_channels=3, num_classes=4) + +TESTS = [] +for type in (VanillaGrad, SmoothGrad, GuidedBackpropGrad, GuidedBackpropSmoothGrad): + # 2D densenet + TESTS.append([type, DENSENET2D, (1, 1, 48, 64), (1, 1, 48, 64)]) + # 3D densenet + TESTS.append([type, DENSENET3D, (1, 1, 6, 6, 6), (1, 1, 6, 6, 6)]) + # 2D senet + TESTS.append([type, SENET2D, (1, 3, 64, 64), (1, 1, 64, 64)]) + # 3D senet + TESTS.append([type, SENET3D, (1, 3, 8, 8, 48), (1, 1, 8, 8, 48)]) + + +class TestGradientClassActivationMap(unittest.TestCase): + @parameterized.expand(TESTS) + def test_shape(self, vis_type, model, shape, expected_shape): + device = "cuda:0" if torch.cuda.is_available() else "cpu" + model.to(device) + model.eval() + vis = vis_type(model) + x = torch.rand(shape, device=device) + result = vis(x) + self.assertTupleEqual(result.shape, x.shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_vis_gradcam.py b/tests/test_vis_gradcam.py index acca06d405..755f4d49ae 100644 --- a/tests/test_vis_gradcam.py +++ b/tests/test_vis_gradcam.py @@ -85,7 +85,7 @@ def test_ill(self): x.requires_grad = False cam = GradCAM(nn_module=model, target_layers="class_layers.relu") image = torch.rand((2, 1, 48, 64)) - with self.assertRaises(RuntimeError): + with self.assertRaises(IndexError): cam(x=image) diff --git a/tests/test_warp.py b/tests/test_warp.py index c039b57211..31f3540c9e 100644 --- a/tests/test_warp.py +++ b/tests/test_warp.py @@ -153,8 +153,9 @@ def test_grad(self): def load_img_and_sample_ddf(): # load image img = LoadImaged(keys="img")({"img": FILE_PATH})["img"] + img = img.detach().numpy() # W, H, D -> D, H, W - img = img.transpose((2, 1, 0)) + img = img.transpose((2, 1, 0)).copy() # randomly sample ddf such that maximum displacement in each direction equals to one-tenth of the image dimension in # that direction. diff --git a/tests/test_weight_init.py b/tests/test_weight_init.py new file mode 100644 index 0000000000..c850ff4ce6 --- /dev/null +++ b/tests/test_weight_init.py @@ -0,0 +1,47 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks.layers import trunc_normal_ + +TEST_CASES = [ + [{"mean": 0.0, "std": 1.0, "a": 2, "b": 4}, (6, 12, 3, 1, 7)], + [{"mean": 0.3, "std": 0.6, "a": -1.0, "b": 1.3}, (1, 4, 4, 4)], + [{"mean": 0.1, "std": 0.4, "a": 1.3, "b": 1.8}, (5, 7, 7, 8, 9)], +] + +TEST_ERRORS = [ + [{"mean": 0.0, "std": 1.0, "a": 5, "b": 1.1}, (1, 1, 8, 8, 8)], + [{"mean": 0.3, "std": -0.1, "a": 1.0, "b": 2.0}, (8, 5, 2, 6, 9)], + [{"mean": 0.7, "std": 0.0, "a": 0.1, "b": 2.0}, (4, 12, 23, 17)], +] + + +class TestWeightInit(unittest.TestCase): + @parameterized.expand(TEST_CASES) + def test_shape(self, input_param, input_shape): + im = torch.rand(input_shape) + trunc_normal_(im, **input_param) + self.assertEqual(im.shape, input_shape) + + @parameterized.expand(TEST_ERRORS) + def test_ill_arg(self, input_param, input_shape): + with self.assertRaises(ValueError): + im = torch.rand(input_shape) + trunc_normal_(im, **input_param) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_wsireader.py b/tests/test_wsireader.py index 6ee02143b8..a0a076b682 100644 --- a/tests/test_wsireader.py +++ b/tests/test_wsireader.py @@ -20,7 +20,7 @@ from monai.data import DataLoader, Dataset from monai.data.image_reader import WSIReader -from monai.transforms import Compose, LoadImaged, ToTensord +from monai.transforms import Compose, FromMetaTensord, LoadImaged, ToTensord from monai.utils import first, optional_import from monai.utils.enums import PostFix from tests.utils import download_url_or_skip_test, testing_data_config @@ -84,7 +84,9 @@ TEST_CASE_RGB_1 = [np.ones((3, 100, 100), dtype=np.uint8)] # CHW -TEST_CASE_ERROR_GRAY = [np.ones((16, 16), dtype=np.uint8)] # no color channel +TEST_CASE_ERROR_0C = [np.ones((16, 16), dtype=np.uint8)] # no color channel +TEST_CASE_ERROR_1C = [np.ones((16, 16, 1), dtype=np.uint8)] # one color channel +TEST_CASE_ERROR_2C = [np.ones((16, 16, 2), dtype=np.uint8)] # two color channels TEST_CASE_ERROR_3D = [np.ones((16, 16, 16, 3), dtype=np.uint8)] # 3D + color @@ -106,22 +108,8 @@ def save_rgba_tiff(array: np.ndarray, filename: str, mode: str): return filename -def save_gray_tiff(array: np.ndarray, filename: str): - """ - Save numpy array into a TIFF file - - Args: - array: numpy ndarray with any shape - filename: the filename to be used for the tiff file. - """ - img_gray = array - imwrite(filename, img_gray, shape=img_gray.shape, photometric="rgb") - - return filename - - @skipUnless(has_cucim or has_osl or has_tiff, "Requires cucim, openslide, or tifffile!") -def setUpModule(): # noqa: N802 +def setUpModule(): hash_type = testing_data_config("images", FILE_KEY, "hash_type") hash_val = testing_data_config("images", FILE_KEY, "hash_val") download_url_or_skip_test(FILE_URL, FILE_PATH, hash_type=hash_type, hash_val=hash_val) @@ -187,13 +175,15 @@ def test_read_rgba(self, img_expected): self.assertIsNone(assert_array_equal(image["RGB"], img_expected)) self.assertIsNone(assert_array_equal(image["RGBA"], img_expected)) - @parameterized.expand([TEST_CASE_ERROR_GRAY, TEST_CASE_ERROR_3D]) + @parameterized.expand([TEST_CASE_ERROR_0C, TEST_CASE_ERROR_1C, TEST_CASE_ERROR_2C, TEST_CASE_ERROR_3D]) @skipUnless(has_tiff, "Requires tifffile.") def test_read_malformats(self, img_expected): + if self.backend == "cucim" and (len(img_expected.shape) < 3 or img_expected.shape[2] == 1): + # Until cuCIM addresses https://github.com/rapidsai/cucim/issues/230 + return reader = WSIReader(self.backend) - file_path = save_gray_tiff( - img_expected, os.path.join(os.path.dirname(__file__), "testing_data", "temp_tiff_image_gray.tiff") - ) + file_path = os.path.join(os.path.dirname(__file__), "testing_data", "temp_tiff_image_gray.tiff") + imwrite(file_path, img_expected, shape=img_expected.shape) with self.assertRaises((RuntimeError, ValueError, openslide.OpenSlideError if has_osl else ValueError)): with reader.read(file_path) as img_obj: reader.get_data(img_obj) @@ -203,6 +193,7 @@ def test_with_dataloader(self, file_path, level, expected_spatial_shape, expecte train_transform = Compose( [ LoadImaged(keys=["image"], reader=WSIReader, backend=self.backend, level=level), + FromMetaTensord(keys=["image"]), ToTensord(keys=["image"]), ] ) diff --git a/tests/test_wsireader_new.py b/tests/test_wsireader_new.py new file mode 100644 index 0000000000..0d5e5892e6 --- /dev/null +++ b/tests/test_wsireader_new.py @@ -0,0 +1,277 @@ +# 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. + +import os +import unittest +from unittest import skipUnless + +import numpy as np +from numpy.testing import assert_array_equal +from parameterized import parameterized + +from monai.data import DataLoader, Dataset +from monai.data.wsi_reader import WSIReader +from monai.transforms import Compose, FromMetaTensord, LoadImaged, ToTensord +from monai.utils import first, optional_import +from monai.utils.enums import PostFix +from tests.utils import assert_allclose, download_url_or_skip_test, testing_data_config + +cucim, has_cucim = optional_import("cucim") +has_cucim = has_cucim and hasattr(cucim, "CuImage") +openslide, has_osl = optional_import("openslide") +imwrite, has_tiff = optional_import("tifffile", name="imwrite") +_, has_codec = optional_import("imagecodecs") +has_tiff = has_tiff and has_codec + +FILE_KEY = "wsi_img" +FILE_URL = testing_data_config("images", FILE_KEY, "url") +base_name, extension = os.path.basename(f"{FILE_URL}"), ".tiff" +FILE_PATH = os.path.join(os.path.dirname(__file__), "testing_data", "temp_" + base_name + extension) + +HEIGHT = 32914 +WIDTH = 46000 + +TEST_CASE_0 = [FILE_PATH, 2, (3, HEIGHT // 4, WIDTH // 4)] + +TEST_CASE_TRANSFORM_0 = [FILE_PATH, 4, (HEIGHT // 16, WIDTH // 16), (1, 3, HEIGHT // 16, WIDTH // 16)] + +TEST_CASE_1 = [ + FILE_PATH, + {}, + {"location": (HEIGHT // 2, WIDTH // 2), "size": (2, 1), "level": 0}, + np.array([[[246], [246]], [[246], [246]], [[246], [246]]]), +] + +TEST_CASE_2 = [ + FILE_PATH, + {}, + {"location": (0, 0), "size": (2, 1), "level": 2}, + np.array([[[239], [239]], [[239], [239]], [[239], [239]]]), +] + +TEST_CASE_3 = [ + FILE_PATH, + {"channel_dim": -1}, + {"location": (HEIGHT // 2, WIDTH // 2), "size": (2, 1), "level": 0}, + np.moveaxis(np.array([[[246], [246]], [[246], [246]], [[246], [246]]]), 0, -1), +] + +TEST_CASE_4 = [ + FILE_PATH, + {"channel_dim": 2}, + {"location": (0, 0), "size": (2, 1), "level": 2}, + np.moveaxis(np.array([[[239], [239]], [[239], [239]], [[239], [239]]]), 0, -1), +] + +TEST_CASE_MULTI_WSI = [ + [FILE_PATH, FILE_PATH], + {"location": (0, 0), "size": (2, 1), "level": 2}, + np.concatenate( + [ + np.array([[[239], [239]], [[239], [239]], [[239], [239]]]), + np.array([[[239], [239]], [[239], [239]], [[239], [239]]]), + ], + axis=0, + ), +] + + +TEST_CASE_RGB_0 = [np.ones((3, 2, 2), dtype=np.uint8)] # CHW + +TEST_CASE_RGB_1 = [np.ones((3, 100, 100), dtype=np.uint8)] # CHW + +TEST_CASE_ERROR_0C = [np.ones((16, 16), dtype=np.uint8)] # no color channel +TEST_CASE_ERROR_1C = [np.ones((16, 16, 1), dtype=np.uint8)] # one color channel +TEST_CASE_ERROR_2C = [np.ones((16, 16, 2), dtype=np.uint8)] # two color channels +TEST_CASE_ERROR_3D = [np.ones((16, 16, 16, 3), dtype=np.uint8)] # 3D + color + + +def save_rgba_tiff(array: np.ndarray, filename: str, mode: str): + """ + Save numpy array into a TIFF RGB/RGBA file + + Args: + array: numpy ndarray with the shape of CxHxW and C==3 representing a RGB image + filename: the filename to be used for the tiff file. '_RGB.tiff' or '_RGBA.tiff' will be appended to this filename. + mode: RGB or RGBA + """ + if mode == "RGBA": + array = np.concatenate([array, 255 * np.ones_like(array[0])[np.newaxis]]).astype(np.uint8) + + img_rgb = array.transpose(1, 2, 0) + imwrite(filename, img_rgb, shape=img_rgb.shape, tile=(16, 16)) + + return filename + + +def save_gray_tiff(array: np.ndarray, filename: str): + """ + Save numpy array into a TIFF file + + Args: + array: numpy ndarray with any shape + filename: the filename to be used for the tiff file. + """ + img_gray = array + imwrite(filename, img_gray, shape=img_gray.shape) + + return filename + + +@skipUnless(has_cucim or has_osl or has_tiff, "Requires cucim, openslide, or tifffile!") +def setUpModule(): + hash_type = testing_data_config("images", FILE_KEY, "hash_type") + hash_val = testing_data_config("images", FILE_KEY, "hash_val") + download_url_or_skip_test(FILE_URL, FILE_PATH, hash_type=hash_type, hash_val=hash_val) + + +class WSIReaderTests: + class Tests(unittest.TestCase): + backend = None + + @parameterized.expand([TEST_CASE_0]) + def test_read_whole_image(self, file_path, level, expected_shape): + reader = WSIReader(self.backend, level=level) + with reader.read(file_path) as img_obj: + img, meta = reader.get_data(img_obj) + self.assertTupleEqual(img.shape, expected_shape) + self.assertEqual(meta["backend"], self.backend) + self.assertEqual(meta["path"], str(os.path.abspath(file_path))) + self.assertEqual(meta["patch_level"], level) + assert_array_equal(meta["patch_size"], expected_shape[1:]) + assert_array_equal(meta["patch_location"], (0, 0)) + + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) + def test_read_region(self, file_path, kwargs, patch_info, expected_img): + reader = WSIReader(self.backend, **kwargs) + with reader.read(file_path) as img_obj: + if self.backend == "tifffile": + with self.assertRaises(ValueError): + reader.get_data(img_obj, **patch_info)[0] + else: + # Read twice to check multiple calls + img, meta = reader.get_data(img_obj, **patch_info) + img2 = reader.get_data(img_obj, **patch_info)[0] + self.assertTupleEqual(img.shape, img2.shape) + self.assertIsNone(assert_array_equal(img, img2)) + self.assertTupleEqual(img.shape, expected_img.shape) + self.assertIsNone(assert_array_equal(img, expected_img)) + self.assertEqual(meta["backend"], self.backend) + self.assertEqual(meta["path"], str(os.path.abspath(file_path))) + self.assertEqual(meta["patch_level"], patch_info["level"]) + assert_array_equal(meta["patch_size"], patch_info["size"]) + assert_array_equal(meta["patch_location"], patch_info["location"]) + + @parameterized.expand([TEST_CASE_MULTI_WSI]) + def test_read_region_multi_wsi(self, file_path_list, patch_info, expected_img): + kwargs = {"name": None, "offset": None} if self.backend == "tifffile" else {} + reader = WSIReader(self.backend, **kwargs) + img_obj_list = reader.read(file_path_list, **kwargs) + if self.backend == "tifffile": + with self.assertRaises(ValueError): + reader.get_data(img_obj_list, **patch_info)[0] + else: + # Read twice to check multiple calls + img, meta = reader.get_data(img_obj_list, **patch_info) + img2 = reader.get_data(img_obj_list, **patch_info)[0] + self.assertTupleEqual(img.shape, img2.shape) + self.assertIsNone(assert_array_equal(img, img2)) + self.assertTupleEqual(img.shape, expected_img.shape) + self.assertIsNone(assert_array_equal(img, expected_img)) + self.assertEqual(meta["backend"], self.backend) + self.assertEqual(meta["path"][0], str(os.path.abspath(file_path_list[0]))) + self.assertEqual(meta["patch_level"][0], patch_info["level"]) + assert_array_equal(meta["patch_size"][0], expected_img.shape[1:]) + assert_array_equal(meta["patch_location"][0], patch_info["location"]) + + @parameterized.expand([TEST_CASE_RGB_0, TEST_CASE_RGB_1]) + @skipUnless(has_tiff, "Requires tifffile.") + def test_read_rgba(self, img_expected): + # skip for OpenSlide since not working with images without tiles + if self.backend == "openslide": + return + image = {} + reader = WSIReader(self.backend) + for mode in ["RGB", "RGBA"]: + file_path = save_rgba_tiff( + img_expected, + os.path.join(os.path.dirname(__file__), "testing_data", f"temp_tiff_image_{mode}.tiff"), + mode=mode, + ) + with reader.read(file_path) as img_obj: + image[mode], _ = reader.get_data(img_obj) + + self.assertIsNone(assert_array_equal(image["RGB"], img_expected)) + self.assertIsNone(assert_array_equal(image["RGBA"], img_expected)) + + @parameterized.expand([TEST_CASE_ERROR_0C, TEST_CASE_ERROR_1C, TEST_CASE_ERROR_2C, TEST_CASE_ERROR_3D]) + @skipUnless(has_tiff, "Requires tifffile.") + def test_read_malformats(self, img_expected): + if self.backend == "cucim" and (len(img_expected.shape) < 3 or img_expected.shape[2] == 1): + # Until cuCIM addresses https://github.com/rapidsai/cucim/issues/230 + return + reader = WSIReader(self.backend) + file_path = os.path.join(os.path.dirname(__file__), "testing_data", "temp_tiff_image_gray.tiff") + imwrite(file_path, img_expected, shape=img_expected.shape) + with self.assertRaises((RuntimeError, ValueError, openslide.OpenSlideError if has_osl else ValueError)): + with reader.read(file_path) as img_obj: + reader.get_data(img_obj) + + @parameterized.expand([TEST_CASE_TRANSFORM_0]) + def test_with_dataloader(self, file_path, level, expected_spatial_shape, expected_shape): + train_transform = Compose( + [ + LoadImaged(keys=["image"], reader=WSIReader, backend=self.backend, level=level), + FromMetaTensord(keys=["image"]), + ToTensord(keys=["image"]), + ] + ) + dataset = Dataset([{"image": file_path}], transform=train_transform) + data_loader = DataLoader(dataset) + data: dict = first(data_loader) + for s in data[PostFix.meta("image")]["spatial_shape"]: + assert_allclose(s, expected_spatial_shape, type_test=False) + self.assertTupleEqual(data["image"].shape, expected_shape) + + @parameterized.expand([TEST_CASE_TRANSFORM_0]) + def test_with_dataloader_batch(self, file_path, level, expected_spatial_shape, expected_shape): + train_transform = Compose( + [ + LoadImaged(keys=["image"], reader=WSIReader, backend=self.backend, level=level), + FromMetaTensord(keys=["image"]), + ToTensord(keys=["image"]), + ] + ) + dataset = Dataset([{"image": file_path}, {"image": file_path}], transform=train_transform) + batch_size = 2 + data_loader = DataLoader(dataset, batch_size=batch_size) + data: dict = first(data_loader) + for s in data[PostFix.meta("image")]["spatial_shape"]: + assert_allclose(s, expected_spatial_shape, type_test=False) + self.assertTupleEqual(data["image"].shape, (batch_size, *expected_shape[1:])) + + +@skipUnless(has_cucim, "Requires cucim") +class TestCuCIM(WSIReaderTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "cucim" + + +@skipUnless(has_osl, "Requires openslide") +class TestOpenSlide(WSIReaderTests.Tests): + @classmethod + def setUpClass(cls): + cls.backend = "openslide" + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_zoom.py b/tests/test_zoom.py index 1a7694072e..78beec69a1 100644 --- a/tests/test_zoom.py +++ b/tests/test_zoom.py @@ -16,8 +16,9 @@ from parameterized import parameterized from scipy.ndimage import zoom as zoom_scipy +from monai.data import MetaTensor, set_track_meta from monai.transforms import Zoom -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion VALID_CASES = [(1.5, "nearest"), (1.5, "nearest"), (0.8, "bilinear"), (0.8, "area")] @@ -27,9 +28,11 @@ class TestZoom(NumpyImageTestCase2D): @parameterized.expand(VALID_CASES) def test_correct_results(self, zoom, mode): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: zoom_fn = Zoom(zoom=zoom, mode=mode, keep_size=False) - zoomed = zoom_fn(p(self.imt[0])) + im = p(self.imt[0]) + zoomed = zoom_fn(im) + test_local_inversion(zoom_fn, zoomed, im) _order = 0 if mode.endswith("linear"): _order = 1 @@ -37,27 +40,37 @@ def test_correct_results(self, zoom, mode): for channel in self.imt[0]: expected.append(zoom_scipy(channel, zoom=zoom, mode="nearest", order=_order, prefilter=False)) expected = np.stack(expected).astype(np.float32) - assert_allclose(zoomed, p(expected), atol=1.0) + assert_allclose(zoomed, p(expected), atol=1.0, type_test=False) def test_keep_size(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: zoom_fn = Zoom(zoom=[0.6, 0.6], keep_size=True, align_corners=True) - zoomed = zoom_fn(p(self.imt[0]), mode="bilinear") - assert_allclose(zoomed.shape, self.imt.shape[1:]) + im = p(self.imt[0]) + zoomed = zoom_fn(im, mode="bilinear") + assert_allclose(zoomed.shape, self.imt.shape[1:], type_test=False) + test_local_inversion(zoom_fn, zoomed, im) zoom_fn = Zoom(zoom=[1.3, 1.3], keep_size=True) - zoomed = zoom_fn(p(self.imt[0])) - assert_allclose(zoomed.shape, self.imt.shape[1:]) + im = p(self.imt[0]) + zoomed = zoom_fn(im) + assert_allclose(zoomed.shape, self.imt.shape[1:], type_test=False) + test_local_inversion(zoom_fn, zoomed, p(self.imt[0])) + + set_track_meta(False) + rotated = zoom_fn(im) + self.assertNotIsInstance(rotated, MetaTensor) + np.testing.assert_allclose(zoomed.shape, self.imt.shape[1:]) + set_track_meta(True) @parameterized.expand(INVALID_CASES) def test_invalid_inputs(self, zoom, mode, raises): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: with self.assertRaises(raises): zoom_fn = Zoom(zoom=zoom, mode=mode) zoom_fn(p(self.imt[0])) def test_padding_mode(self): - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: zoom_fn = Zoom(zoom=0.5, mode="nearest", padding_mode="constant", keep_size=True) test_data = p([[[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]]) zoomed = zoom_fn(test_data) diff --git a/tests/test_zoomd.py b/tests/test_zoomd.py index 87a5cec22b..b6ff86e474 100644 --- a/tests/test_zoomd.py +++ b/tests/test_zoomd.py @@ -16,7 +16,7 @@ from scipy.ndimage import zoom as zoom_scipy from monai.transforms import Zoomd -from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose +from tests.utils import TEST_NDARRAYS_ALL, NumpyImageTestCase2D, assert_allclose, test_local_inversion VALID_CASES = [(1.5, "nearest", False), (0.3, "bilinear", False), (0.8, "bilinear", False)] @@ -28,8 +28,10 @@ class TestZoomd(NumpyImageTestCase2D): def test_correct_results(self, zoom, mode, keep_size): key = "img" zoom_fn = Zoomd(key, zoom=zoom, mode=mode, keep_size=keep_size) - for p in TEST_NDARRAYS: - zoomed = zoom_fn({key: p(self.imt[0])}) + for p in TEST_NDARRAYS_ALL: + im = p(self.imt[0]) + zoomed = zoom_fn({key: im}) + test_local_inversion(zoom_fn, zoomed, {key: im}, key) _order = 0 if mode.endswith("linear"): _order = 1 @@ -38,12 +40,12 @@ def test_correct_results(self, zoom, mode, keep_size): ] expected = np.stack(expected).astype(np.float32) - assert_allclose(zoomed[key], p(expected), atol=1.0) + assert_allclose(zoomed[key], p(expected), atol=1.0, type_test=False) def test_keep_size(self): key = "img" zoom_fn = Zoomd(key, zoom=0.6, keep_size=True, padding_mode="constant", constant_values=2) - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: zoomed = zoom_fn({key: p(self.imt[0])}) np.testing.assert_array_equal(zoomed[key].shape, self.imt.shape[1:]) @@ -54,7 +56,7 @@ def test_keep_size(self): @parameterized.expand(INVALID_CASES) def test_invalid_inputs(self, _, zoom, mode, raises): key = "img" - for p in TEST_NDARRAYS: + for p in TEST_NDARRAYS_ALL: with self.assertRaises(raises): zoom_fn = Zoomd(key, zoom=zoom, mode=mode) zoom_fn({key: p(self.imt[0])}) diff --git a/tests/testing_data/CT_DICOM/17106 b/tests/testing_data/CT_DICOM/7106 similarity index 92% rename from tests/testing_data/CT_DICOM/17106 rename to tests/testing_data/CT_DICOM/7106 index 34c9659147..727bea124b 100644 Binary files a/tests/testing_data/CT_DICOM/17106 and b/tests/testing_data/CT_DICOM/7106 differ diff --git a/tests/testing_data/CT_DICOM/17136 b/tests/testing_data/CT_DICOM/7136 similarity index 92% rename from tests/testing_data/CT_DICOM/17136 rename to tests/testing_data/CT_DICOM/7136 index 81949d1077..ed1222f80d 100644 Binary files a/tests/testing_data/CT_DICOM/17136 and b/tests/testing_data/CT_DICOM/7136 differ diff --git a/tests/testing_data/CT_DICOM/17166 b/tests/testing_data/CT_DICOM/7166 similarity index 92% rename from tests/testing_data/CT_DICOM/17166 rename to tests/testing_data/CT_DICOM/7166 index e9cbc38b01..edadf552b2 100644 Binary files a/tests/testing_data/CT_DICOM/17166 and b/tests/testing_data/CT_DICOM/7166 differ diff --git a/tests/testing_data/CT_DICOM/17196 b/tests/testing_data/CT_DICOM/7196 similarity index 92% rename from tests/testing_data/CT_DICOM/17196 rename to tests/testing_data/CT_DICOM/7196 index cc579d5425..d4c0f0795e 100644 Binary files a/tests/testing_data/CT_DICOM/17196 and b/tests/testing_data/CT_DICOM/7196 differ diff --git a/tests/testing_data/data_config.json b/tests/testing_data/data_config.json index e337a2d2e7..254314d1b8 100644 --- a/tests/testing_data/data_config.json +++ b/tests/testing_data/data_config.json @@ -34,6 +34,26 @@ "url": "https://github.com/rcuocolo/PROSTATEx_masks/raw/master/Files/lesions/Images/ADC/ProstateX-0000_ep2d_diff_tra_7.nii.gz", "hash_type": "md5", "hash_val": "f12a11ad0ebb0b1876e9e010564745d2" + }, + "ref_avg152T1_LR": { + "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/avg152T1_LR_nifti.nii.gz", + "hash_type": "sha256", + "hash_val": "c01a50caa7a563158ecda43d93a1466bfc8aa939bc16b06452ac1089c54661c8" + }, + "ref_avg152T1_RL": { + "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/avg152T1_RL_nifti.nii.gz", + "hash_type": "sha256", + "hash_val": "8a731128dac4de46ccb2cc60d972b98f75a52f21fb63ddb040ca96f0aed8b51a" + }, + "MNI152_T1_2mm": { + "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/MNI152_T1_2mm.nii.gz", + "hash_type": "sha256", + "hash_val": "0585cd056bf5ccfb8bf97a5f6a66082d4e7caad525718fc11e40d80a827fcb92" + }, + "MNI152_T1_2mm_strucseg": { + "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/MNI152_T1_2mm_strucseg.nii.gz", + "hash_type": "sha256", + "hash_val": "eb4f1e596ca85aadaefc359d409fb9a3e27d733e6def04b996953b7c54bc26d4" } }, "models": { @@ -72,7 +92,7 @@ "test_meta_file": { "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json", "hash_type": "md5", - "hash_val": "e12813de2c15672a8c8fa8466b3dfc95" + "hash_val": "662135097106b71067cd1fc657f8720f" } } } diff --git a/tests/testing_data/inference.json b/tests/testing_data/inference.json index cc9ddef866..46aa206d03 100644 --- a/tests/testing_data/inference.json +++ b/tests/testing_data/inference.json @@ -1,8 +1,8 @@ { "dataset_dir": "/workspace/data/Task09_Spleen", + "prediction_shape": "prediction shape:", "import_glob": "$import glob", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", - "set_seed": "$monai.utils.set_determinism(0)", "print_test_name": "$print('json_test')", "print_glob_file": "$print(glob.__file__)", "network_def": { @@ -46,10 +46,6 @@ "_target_": "RandRotated", "_disabled_": true, "keys": "image" - }, - { - "_target_": "EnsureTyped", - "keys": "image" } ] }, @@ -91,15 +87,19 @@ { "_target_": "SaveImaged", "keys": "pred", - "meta_keys": "image_meta_dict", "output_dir": "@_meta_#output_dir" + }, + { + "_target_": "Lambdad", + "keys": "pred", + "func": "$lambda x: print(@prediction_shape + str(x.shape))", + "overwrite": false } ] }, "evaluator": { "_target_": "SupervisedEvaluator", "_requires_": [ - "@set_seed", "@print_test_name", "@print_glob_file", "$print('test_in_line_json')" @@ -110,5 +110,9 @@ "inferer": "@inferer", "postprocessing": "@postprocessing", "amp": false - } + }, + "evaluating": [ + "$monai.utils.set_determinism(0)", + "$@evaluator.run()" + ] } diff --git a/tests/testing_data/inference.yaml b/tests/testing_data/inference.yaml index 4973d4473f..90f0bb35b9 100644 --- a/tests/testing_data/inference.yaml +++ b/tests/testing_data/inference.yaml @@ -1,7 +1,7 @@ --- dataset_dir: "/workspace/data/Task09_Spleen" +prediction_shape: "prediction shape:" device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')" -set_seed: "$monai.utils.set_determinism(0)" print_test_name: "$print('yaml_test')" network_def: _target_: UNet @@ -34,8 +34,6 @@ preprocessing: - _target_: RandRotated _disabled_: true keys: image - - _target_: EnsureTyped - keys: image dataset: _target_: need override data: "@_meta_#datalist" @@ -65,13 +63,15 @@ postprocessing: argmax: true - _target_: SaveImaged keys: pred - meta_keys: image_meta_dict output_dir: "@_meta_#output_dir" + - _target_: Lambdad + keys: pred + func: "$lambda x: print(@prediction_shape + str(x.shape))" + overwrite: false evaluator: _target_: SupervisedEvaluator _requires_: - "$print('test_in_line_yaml')" - - "@set_seed" - "@print_test_name" device: "@device" val_data_loader: "@dataloader" @@ -79,3 +79,6 @@ evaluator: inferer: "@inferer" postprocessing: "@postprocessing" amp: false +evaluating: + - "$monai.utils.set_determinism(0)" + - "$@evaluator.run()" diff --git a/tests/testing_data/matshow3d_patch_test.png b/tests/testing_data/matshow3d_patch_test.png index a4d89e3446..0a4632a763 100644 Binary files a/tests/testing_data/matshow3d_patch_test.png and b/tests/testing_data/matshow3d_patch_test.png differ diff --git a/tests/testing_data/metadata.json b/tests/testing_data/metadata.json index c12ae411f1..98a17b73c5 100644 --- a/tests/testing_data/metadata.json +++ b/tests/testing_data/metadata.json @@ -5,7 +5,7 @@ "0.1.0": "complete the model package", "0.0.1": "initialize the model package structure" }, - "monai_version": "0.8.0", + "monai_version": "0.9.0", "pytorch_version": "1.10.0", "numpy_version": "1.21.2", "optional_packages_version": { diff --git a/tests/testing_data/transform_metatensor_cases.yaml b/tests/testing_data/transform_metatensor_cases.yaml new file mode 100644 index 0000000000..75275b9df5 --- /dev/null +++ b/tests/testing_data/transform_metatensor_cases.yaml @@ -0,0 +1,197 @@ +--- +input_keys: [image, segs] +test_device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')" +init_affine: "$np.array([[-2, 0, 0, 90], [0, 2, 0, -126], [0, 0, 2, -72], [0, 0, 0, 1]], dtype=np.float64)" +init_shape: [1, 91, 109, 91] +TEST_CASE_1: + _target_: Compose + transforms: + - _target_: LoadImageD + keys: "@input_keys" + ensure_channel_first: True + image_only: True + - _target_: ToDeviced + keys: "@input_keys" + device: "@test_device" + - _target_: CenterScaleCropD + keys: "@input_keys" + roi_scale: 0.98 + - _target_: CropForegroundD + keys: "@input_keys" + source_key: seg + start_coord_key: null + end_coord_key: null + k_divisible: 5 + - _target_: RandSpatialCropD + keys: "@input_keys" + roi_size: [76, 87, 73] + - _target_: RandScaleCropD + keys: "@input_keys" + roi_scale: 0.9 + - _target_: ResizeWithPadOrCropD + keys: "@input_keys" + spatial_size: [32, 43, 54] + - _target_: DivisiblePadD + keys: "@input_keys" + k: 3 + +TEST_CASE_2: + _target_: Compose + transforms: + - _target_: LoadImaged + keys: "@input_keys" + ensure_channel_first: False + image_only: True + - _target_: ToDeviced + keys: "@input_keys" + device: "@test_device" + - _target_: EnsureChannelFirstd + keys: "@input_keys" + - _target_: ScaleIntensityRangePercentilesd + keys: "$@input_keys[0]" + lower: 4 + upper: 95 + b_min: 1 + b_max: 10 + - _target_: RandScaleIntensityd + keys: "$@input_keys[0]" + prob: 1.0 + factors: [5, 10] + - _target_: RandGaussianNoised + keys: "$@input_keys[0]" + prob: 1.0 + mean: 10.0 + std: 2.0 + - _target_: RandCoarseShuffled + keys: "$@input_keys[0]" + prob: 1.0 + holes: 2 + spatial_size: [10, 13, 18] + max_spatial_size: [14, 30, 57] + - _target_: DataStatsd + keys: "$@input_keys[0]" + - _target_: RandBiasFieldd + keys: "$@input_keys[0]" + prob: 1.0 + - _target_: RandGaussianSmoothd + keys: "$@input_keys[0]" + prob: 1.0 + - _target_: RandGaussianSharpend + keys: "$@input_keys[0]" + prob: 1.0 + - _target_: RandHistogramShiftd + keys: "$@input_keys[0]" + prob: 1.0 + - _target_: RandAdjustContrastd + keys: "$@input_keys[0]" + prob: 1.0 + - _target_: RandCoarseDropoutd + keys: "$@input_keys[0]" + prob: 1.0 + holes: 3 + spatial_size: [10, 13, 18] + max_spatial_size: [14, 30, 57] + - _target_: RandRicianNoised + keys: "$@input_keys[0]" + prob: 1.0 + +TEST_CASE_3: + _target_: Compose + transforms: + - _target_: LoadImageD + keys: "@input_keys" + ensure_channel_first: True + image_only: True + - _target_: CenterScaleCropD + keys: "@input_keys" + roi_scale: 0.98 + - _target_: CropForegroundD + keys: "@input_keys" + source_key: seg + start_coord_key: null + end_coord_key: null + k_divisible: 5 + - _target_: ToDeviced + keys: "@input_keys" + device: "@test_device" + - _target_: RandRotate90d + keys: "@input_keys" + prob: 1.0 + spatial_axes: [2, 1] + - _target_: Spacingd + keys: "@input_keys" + pixdim: [1.8, 2.1, 2.3] + - _target_: RandFlipd + keys: "@input_keys" + prob: 1.0 + spatial_axis: 2 + - _target_: RandAffined + keys: "@input_keys" + prob: 1.0 + spatial_size: [80, 91, 92] + rotate_range: 1.0 + scale_range: 0.1 + - _target_: Flipd + keys: "@input_keys" + spatial_axis: 2 + - _target_: Orientationd + keys: "@input_keys" + axcodes: "RPI" + - _target_: Affined + keys: "@input_keys" + shear_params: [0, 0.5, 0] + - _target_: Rotate90d + keys: "@input_keys" + spatial_axes: [1, 2] + - _target_: Zoomd + keys: "@input_keys" + zoom: 1.3 + - _target_: ScaleIntensityd + keys: "@input_keys" + minv: 0 + maxv: 10 + - _target_: RandAxisFlipD + keys: "@input_keys" + prob: 1.0 + - _target_: RandRotated + keys: "@input_keys" + prob: 1.0 + range_y: "$np.pi/3" + - _target_: RandZoomD + keys: "@input_keys" + prob: 1.0 + max_zoom: 1.2 + keep_size: True + - _target_: RandGaussianNoised + keys: "@input_keys" + prob: 1.0 + - _target_: ResizeWithPadOrCropD + keys: "@input_keys" + spatial_size: [71, 56, 80] + - _target_: Rand3DElasticd + keys: "@input_keys" + spatial_size: [71, 56, 80] + sigma_range: [5, 7] + magnitude_range: [50, 150] + prob: 1.0 + - _target_: Resized + keys: "@input_keys" + spatial_size: [72, 57, 82] + +TEST_CASE_1_answer: + load_shape: [1, 1, 33, 45, 54] + affine: "$np.array([[-2, 0, 0, 34], [0, 2, 0, -64], [0, 0, 2, -54], [0, 0, 0, 1]], dtype=np.float64)" + inv_affine: "@init_affine" + inv_shape: "@init_shape" + +TEST_CASE_2_answer: + load_shape: [1, 1, 91, 109, 91] + affine: "$np.array([[-2, 0, 0, 90], [0, 2, 0, -126], [0, 0, 2, -72], [0, 0, 0, 1]], dtype=np.float64)" + inv_affine: "@init_affine" + inv_shape: "@init_shape" + +TEST_CASE_3_answer: + load_shape: [1, 1, 72, 57, 82] + affine: "$np.array([[-1.343816, -0.682904, -0.234832, 76.01494], [0.309004, 0.653211, -1.734872, 24.511358], [-0.104049, 1.617199, 0.584171, -56.521294], [0, 0, 0, 1]], dtype=np.float64)" + inv_affine: "@init_affine" + inv_shape: "@init_shape" diff --git a/tests/utils.py b/tests/utils.py index 3065f9b3df..e3e77a9c32 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -26,7 +26,7 @@ from contextlib import contextmanager from functools import partial, reduce from subprocess import PIPE, Popen -from typing import Callable, Optional, Tuple +from typing import Callable, Optional, Tuple, Union from urllib.error import ContentTooShortError, HTTPError import numpy as np @@ -38,6 +38,7 @@ from monai.config.deviceconfig import USE_COMPILED from monai.config.type_definitions import NdarrayOrTensor from monai.data import create_test_image_2d, create_test_image_3d +from monai.data.meta_tensor import MetaTensor, get_track_meta from monai.networks import convert_to_torchscript from monai.utils import optional_import from monai.utils.module import pytorch_after, version_leq @@ -76,7 +77,7 @@ def clone(data: NdarrayTensor) -> NdarrayTensor: def assert_allclose( actual: NdarrayOrTensor, desired: NdarrayOrTensor, - type_test: bool = True, + type_test: Union[bool, str] = True, device_test: bool = False, *args, **kwargs, @@ -88,13 +89,22 @@ def assert_allclose( actual: Pytorch Tensor or numpy array for comparison. desired: Pytorch Tensor or numpy array to compare against. type_test: whether to test that `actual` and `desired` are both numpy arrays or torch tensors. + if type_test == "tensor", it checks whether the `actual` is a torch.tensor or metatensor according to + `get_track_meta`. device_test: whether to test the device property. args: extra arguments to pass on to `np.testing.assert_allclose`. kwargs: extra arguments to pass on to `np.testing.assert_allclose`. """ - if type_test: + if isinstance(type_test, str) and type_test == "tensor": + if get_track_meta(): + np.testing.assert_equal(isinstance(actual, MetaTensor), True, "must be a MetaTensor") + else: + np.testing.assert_equal( + isinstance(actual, torch.Tensor) and not isinstance(actual, MetaTensor), True, "must be a torch.Tensor" + ) + elif type_test: # check both actual and desired are of the same type np.testing.assert_equal(isinstance(actual, np.ndarray), isinstance(desired, np.ndarray), "numpy type") np.testing.assert_equal(isinstance(actual, torch.Tensor), isinstance(desired, torch.Tensor), "torch type") @@ -102,8 +112,8 @@ def assert_allclose( if isinstance(desired, torch.Tensor) or isinstance(actual, torch.Tensor): if device_test: np.testing.assert_equal(str(actual.device), str(desired.device), "torch device check") # type: ignore - actual = actual.cpu().numpy() if isinstance(actual, torch.Tensor) else actual - desired = desired.cpu().numpy() if isinstance(desired, torch.Tensor) else desired + actual = actual.detach().cpu().numpy() if isinstance(actual, torch.Tensor) else actual + desired = desired.detach().cpu().numpy() if isinstance(desired, torch.Tensor) else desired np.testing.assert_allclose(actual, desired, *args, **kwargs) @@ -586,13 +596,11 @@ def _wrapper(*args, **kwargs): def _cache_original_func(obj) -> None: """cache the original function by name, so that the decorator doesn't shadow it.""" - global _original_funcs _original_funcs[obj.__name__] = obj def _del_original_func(obj): """pop the original function from cache.""" - global _original_funcs _original_funcs.pop(obj.__name__, None) if torch.cuda.is_available(): # clean up the cached function torch.cuda.synchronize() @@ -709,10 +717,39 @@ def query_memory(n=2): return ",".join(f"{int(x)}" for x in ids) -TEST_NDARRAYS: Tuple[Callable] = (np.array, torch.as_tensor) # type: ignore +def test_local_inversion(invertible_xform, to_invert, im, dict_key=None): + """test that invertible_xform can bring to_invert back to im""" + im_item = im if dict_key is None else im[dict_key] + if not isinstance(im_item, MetaTensor): + return + im_inv = invertible_xform.inverse(to_invert) + if dict_key: + im_inv = im_inv[dict_key] + im = im[dict_key] + np.testing.assert_array_equal(im_inv.applied_operations, []) + assert_allclose(im_inv.shape, im.shape) + assert_allclose(im_inv.affine, im.affine, atol=1e-3, rtol=1e-3) + + +TEST_TORCH_TENSORS: Tuple[Callable] = (torch.as_tensor,) if torch.cuda.is_available(): gpu_tensor: Callable = partial(torch.as_tensor, device="cuda") - TEST_NDARRAYS = TEST_NDARRAYS + (gpu_tensor,) # type: ignore + TEST_NDARRAYS = TEST_TORCH_TENSORS + (gpu_tensor,) + +DEFAULT_TEST_AFFINE = torch.tensor( + [[2.0, 0.0, 0.0, 0.0], [0.0, 2.0, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0], [0.0, 0.0, 0.0, 1.0]] +) +_metatensor_creator = partial(MetaTensor, meta={"a": "b", "affine": DEFAULT_TEST_AFFINE}) +TEST_NDARRAYS_NO_META_TENSOR: Tuple[Callable] = (np.array,) + TEST_TORCH_TENSORS # type: ignore +TEST_NDARRAYS: Tuple[Callable] = TEST_NDARRAYS_NO_META_TENSOR + (_metatensor_creator,) # type: ignore +TEST_TORCH_AND_META_TENSORS: Tuple[Callable] = TEST_TORCH_TENSORS + (_metatensor_creator,) # type: ignore +# alias for branch tests +TEST_NDARRAYS_ALL = TEST_NDARRAYS + + +TEST_DEVICES = [[torch.device("cpu")]] +if torch.cuda.is_available(): + TEST_DEVICES.append([torch.device("cuda")]) if __name__ == "__main__":