ENH: Add notebook for MONAI with itk-elastix

[ntatsisk]

Adds an example of how to combine itk-elastix as a pre-processing registration step for deep learning using MONAI.

Related to the discussion in #126.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[dzenanz]

Looks good, one minor remark.

[thewtex]

@HastingsGreer @aylward thoughts and suggestions appreciated

[aylward]

Nicely done. Great, compact, and complete example. Serves as an example of how other methods can be used to create custom transforms.

[aylward]

LGTM!

[N-Dekker]

For the record, the error at Notebook tests / Run notebooks on treebeard (pull_request) says:

ModuleNotFoundError                       Traceback (most recent call last)
Cell In [1], line 3
      1 import os
      2 import itk
----> 3 import torch
      4 import numpy as np
      5 from tqdm import tqdm

ModuleNotFoundError: No module named 'torch'

[thewtex]

Line #14.    from monai.apps import MedNISTDataset

Could a cell at the top please be added to install dependencies via pip? This improves reproduciblity and portability. We want the CI tests to pass.

Reply via ReviewNB

[thewtex]

Until we get to that point, we should install the dependencies here: https://github.com/InsightSoftwareConsortium/ITKElastix/blob/main/.binder/requirements.txt

[thewtex]

Very nicely done!

[aylward]

Also - these should probably be invertible transforms. You'll need to store the transform determined and used by elastix as a local var in the transform, and provide a method for applying its inverse. For info, see https://docs.monai.io/en/stable/_modules/monai/transforms/inverse.html
For an example, see Rotate90
https://docs.monai.io/en/stable/_modules/monai/transforms/spatial/array.html#Rotate90

[ntatsisk]

Thanks everyone for the suggestions!

@thewtex, @N-Dekker I added pytorch and MONAI in.binder/requirements.txt.

@dzenanz I agree that aesthetics play an important role. The interpolation was bilinear by default so I changed it to bicubic now, but I am not sure if it affected the final result. What made the final image prettier was to increase the learning rate so that the model converges more. Still, there are some discontinuities left though.

@aylward Great suggestion! However, currently we don't have yet a way to store elastix transforms in memory - we are working on it. I think that the other option would be to store this information in the affine_matrix of each MetaTensor, but I feel that deserves a separate notebook by itself (I am planning to create one). I propose that for the time being we keep these (monai) transforms not-invertible and when we have support for in-memory elastix transforms, we get back to it. What do you think?

[ntatsisk]

ModuleNotFoundError                       Traceback (most recent call last)
Cell In [1], line 5
      3 import torch
      4 import numpy as np
----> 5 from tqdm import tqdm
      6 import matplotlib.pyplot as plt
      8 from monai.utils import set_determinism

ModuleNotFoundError: No module named 'tqdm'

@thewtex, it looks that it now fails because there is no tqdm package. Is it ok if I also add that in the .binder/requirements.txt? It is not mandatory to use but it does make the notebook neater.

[dzenanz]

@ntatsisk please add tqdm as needed.

[ntatsisk]

@thewtex, @dzenanz It seems that all notebook kernels are dead now. Is it a problem of itk-elastix with itk=5.3rc04.post4? Because I see that I have the same issue locally - python exits abruptly without an error and notebook kernel dies again without an error. itk=5.3rc04.post3 works fine though.

[dzenanz]

Could it be InsightSoftwareConsortium/ITKRemoteModuleBuildTestPackageAction#9? Maybe @tbirdso could comment.

[tbirdso]

I am not certain whether the issues are related. I am investigating issues with Windows Python and Linux Python 3.11 builds in InsightSoftwareConsortium/ITKRemoteModuleBuildTestPackageAction#9, but it looks like notebook checks in this PR are getting Linux Python 3.9 wheels which are not known to be affected:

Collecting itk>=5.3rc4.post3
  Downloading itk-5.3rc4.post4-cp39-cp39-manylinux_2_28_x86_64.whl (8.4 kB)

It may be that itk-elastix wheels will need to be rebuilt for ITK v5.3rc04.post4. For now itk==v5.3rc04.post3 could be used for running notebooks.

[tbirdso]

Hi @ntatsisk , ITKElastix 0.14.4 has been released for compatibility with the ITK v5.3rc04.post4 release. Notebook checks are now passing again in the main branch.

Would you please rebase these changes on main and re-run notebook checks?

[ntatsisk]

 RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx
RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

The notebooks seem to run now but this time it gives the above error. It happens when the notebook attempts to move the model to the GPU. Any ideas how to fix this?

[dzenanz]

I think that GitHub actions runner machines don't have GPUs. The usual solution is to make code support CPUs, too. Usual way is described on SO.

[tbirdso]

 RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx
RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

The notebooks seem to run now but this time it gives the above error. It happens when the notebook attempts to move the model to the GPU. Any ideas how to fix this?

Hi @ntatsisk , yes, this is expected as Github CI runners do not have GPUs attached. There are two approaches we can take here:

1. Skip over the notebook for now and do not try to run it as part of CI checks. This can be done with the nbmake ignore parameter, see https://github.com/treebeardtech/nbmake-action/blob/main/.github/workflows/action_unit_test.yml#L54 for a hint as to how to set up the .yml. This is a short-term solution that seems acceptable for now (assuming the notebook runs on your system without issue).
2. Use a self-hosted runner with a GPU attached for notebook tests. We currently do this for accelerated testing in the ITKVkFFTBackend and ITKCLEsperanto repos via CiRun. This is a longer-term solution that will require more setup but would help with maintenance in the long run. @thewtex may have more thoughts on what would be required for setting this up, i.e. an AWS account.

EDIT: I like @dzenanz 's solution best -- if this is a notebook that could run on CPU then it should be amended to do so.

[ntatsisk]

Thank you for the suggestions. For some reason I believed that there was GPU support already here but clearly that is not the case.

I think the decision if this and the future MONAI notebooks run on the CPU or the GPU boils down to what our goal is. Do we aim for complete real-world applications/problems with maybe >2D data or more minimal examples simply showcasing the possibilities with itk-elastix? For example even if the current notebook is made to run on the CPU, I don't think that an upcoming 3D lung segmentation notebook that we discussed with @thewtex will run also on CPU.

If we want to avoid GPU support in this repo, I can see also the solution of having two sets of notebooks: 1) simple ones that are tested (CPU) and 2) more elaborate ones that might use the GPU but are not tested. The latter could live also outside of this repo if it makes more sense. What are your thoughts? (.. cc @N-Dekker, @thewtex )

[thewtex]

@ntatsisk thanks for working on this. My 2 cents on a path forward:

Yes, we should have notebook that runs on CPU and can also utilize GPU.

I think a good path forward is to support CPU for CI for the 2D.

We can set up a GPU CI test in another PR.

We can add the 3D notebook and just make a note that it should be used on a GPU system for a reasonable run time.

[ntatsisk]

Great! I changed the notebook to use the GPU conditionally as @dzenanz pointed out. I also made the number of epochs dependent on whether the notebook runs on CPU or GPU. That way it gets tested in the CPU while the runtime is kept low.
Could you re-trigger the tests?

[tbirdso]

Hi @ntatsisk , is this ready to be merged?

[ntatsisk]

@tbirdso yes! :)