ENH: Add monai MetaTensor support

Applies the .affine matrix.
InsightSoftwareConsortium · Jan 30, 2024 · 0b974ff · 0b974ff
1 parent 68ea8b4
commit 0b974ff
Show file tree

Hide file tree

Showing 3 changed files with 156 additions and 35 deletions.
diff --git a/examples/integrations/MONAI/transform_visualization.ipynb b/examples/integrations/MONAI/transform_visualization.ipynb
diff --git a/itkwidgets/integrations/__init__.py b/itkwidgets/integrations/__init__.py
@@ -6,6 +6,7 @@
 import dask
 from .itk import HAVE_ITK
 from .pytorch import HAVE_TORCH
+from .monai import HAVE_MONAI
 from .vtk import HAVE_VTK, vtk_image_to_ngff_image, vtk_polydata_to_vtkjs
 from .xarray import HAVE_XARRAY, HAVE_MULTISCALE_SPATIAL_IMAGE, xarray_data_array_to_numpy, xarray_data_set_to_numpy
 from ..render_types import RenderType
@@ -102,6 +103,15 @@ def _get_viewer_image(image, label=False):
         to_ngff_zarr(store, multiscales, chunk_store=chunk_store)
         return store
 
+    if HAVE_MONAI:
+        from monai.data import MetaTensor, metatensor_to_itk_image
+        if isinstance(image, MetaTensor):
+            itk_image = metatensor_to_itk_image(image)
+            ngff_image = itk_image_to_ngff_image(itk_image)
+            multiscales = to_multiscales(ngff_image, method=method)
+            to_ngff_zarr(store, multiscales, chunk_store=chunk_store)
+            return store
+
     if HAVE_TORCH:
         import torch
         if isinstance(image, torch.Tensor):

diff --git a/itkwidgets/integrations/monai.py b/itkwidgets/integrations/monai.py
@@ -0,0 +1,8 @@
+import importlib_metadata
+
+HAVE_MONAI = False
+try:
+    importlib_metadata.metadata("monai")
+    HAVE_MONAI = True
+except importlib_metadata.PackageNotFoundError:
+    pass