909 IO factory - develop ITKReader, NibabelReader, LoadImage

examples/notebooks/IO_factory_test.ipynb

@@ -0,0 +1,193 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Multi GPU Test\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Project-MONAI/MONAI/blob/master/examples/notebooks/multi_gpu_test.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install -qU \"monai[itk]\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install -qU itk"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MONAI version: 0.1.0+325.g7de3675.dirty\n",
+      "Python version: 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 19:07:31)  [GCC 7.3.0]\n",
+      "Numpy version: 1.18.1\n",
+      "Pytorch version: 1.6.0\n",
+      "\n",
+      "Optional dependencies:\n",
+      "Pytorch Ignite version: 0.3.0\n",
+      "Nibabel version: 3.1.1\n",
+      "scikit-image version: 0.15.0\n",
+      "Pillow version: 7.0.0\n",
+      "Tensorboard version: 2.1.0\n",
+      "\n",
+      "For details about installing the optional dependencies, please visit:\n",
+      "    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Copyright 2020 MONAI Consortium\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License.\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "from monai.config import print_config\n",
+    "from monai.transforms import LoadImage\n",
+    "\n",
+    "print_config()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test loading Nifti files"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "filename = \"/workspace/data/medical/Task09_Spleen/imagesTr/spleen_10.nii.gz\"\n",
+    "loader = LoadImage()\n",
+    "data, meta = loader(filename)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'ITK_FileNotes': '5.0.10', 'aux_file': '', 'bitpix': '32', 'cal_max': '0', 'cal_min': '0', 'datatype': '16', 'descrip': '5.0.10', 'dim[0]': '3', 'dim[1]': '512', 'dim[2]': '512', 'dim[3]': '55', 'dim[4]': '1', 'dim[5]': '1', 'dim[6]': '1', 'dim[7]': '1', 'dim_info': '0', 'intent_code': '0', 'intent_name': '', 'intent_p1': '0', 'intent_p2': '0', 'intent_p3': '0', 'nifti_type': '1', 'pixdim[0]': '0', 'pixdim[1]': '0.976562', 'pixdim[2]': '0.976562', 'pixdim[3]': '5', 'pixdim[4]': '0', 'pixdim[5]': '0', 'pixdim[6]': '0', 'pixdim[7]': '0', 'qform_code': '1', 'qform_code_name': 'NIFTI_XFORM_SCANNER_ANAT', 'qoffset_x': '-499.023', 'qoffset_y': '-499.023', 'qoffset_z': '0', 'quatern_b': '0', 'quatern_c': '0', 'quatern_d': '0', 'scl_inter': '0', 'scl_slope': '1', 'sform_code': '1', 'sform_code_name': 'NIFTI_XFORM_SCANNER_ANAT', 'slice_code': '0', 'slice_duration': '0', 'slice_end': '0', 'slice_start': '0', 'srow_x': '0.976562 0 0 -499.023', 'srow_y': '0 0.976562 0 -499.023', 'srow_z': '0 0 5 0', 'toffset': '0', 'vox_offset': '352', 'xyzt_units': '10', 'origin': array([499.02319336, 499.02319336,   0.        ]), 'spacing': array([0.97656202, 0.97656202, 5.        ]), 'direction': array([[-1.,  0.,  0.],\n",
+      "       [ 0., -1.,  0.],\n",
+      "       [ 0.,  0.,  1.]]), 'filename_or_obj': '/workspace/data/medical/Task09_Spleen/imagesTr/spleen_10.nii.gz', 'affine': array([[ -0.97656202,   0.        ,   0.        , 499.02319336],\n",
+      "       [  0.        ,  -0.97656202,   0.        , 499.02319336],\n",
+      "       [  0.        ,   0.        ,   5.        ,   0.        ],\n",
+      "       [  0.        ,   0.        ,   0.        ,   1.        ]]), 'original_affine': array([[ -0.97656202,   0.        ,   0.        , 499.02319336],\n",
+      "       [  0.        ,  -0.97656202,   0.        , 499.02319336],\n",
+      "       [  0.        ,   0.        ,   5.        ,   0.        ],\n",
+      "       [  0.        ,   0.        ,   0.        ,   1.        ]]), 'spatial_shape': [55, 512, 512]}\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(meta)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(55, 512, 512)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(data.shape)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.6.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

monai/data/__init__.py

@@ -14,6 +14,7 @@
 from .dataset import ArrayDataset, CacheDataset, Dataset, PersistentDataset, ZipDataset
 from .decathalon_datalist import load_decathalon_datalist
 from .grid_dataset import GridPatchDataset
+from .image_reader import *
 from .nifti_reader import NiftiDataset
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti

monai/data/image_reader.py

@@ -0,0 +1,94 @@
+# Copyright 2020 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 ABC, abstractmethod
+from typing import Any, List, Optional, Sequence, Union
+
+import itk
+import numpy as np
+
+
+class ImageReader(ABC):
+    """Abstract class to define interface APIs to load image files.
+
+    """
+
+    def __init__(self, suffixes: Optional[Union[str, Sequence[str]]] = None, img: Any = None):
+        self.suffixes = suffixes
+        self.img = img
+
+    def verify_suffix(self, suffix: str):
+        return False if self.suffixes is not None and suffix not in self.suffixes else True
+
+    def uncache(self):
+        self.img = None
+
+    @abstractmethod
+    def read_image(self, filename: str):
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+    @abstractmethod
+    def get_meta_dict(self) -> List:
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+    @abstractmethod
+    def get_affine(self) -> List:
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+    @abstractmethod
+    def get_spatial_shape(self) -> List:
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+    @abstractmethod
+    def get_array_data(self) -> np.ndarray:
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+
+class ITKReader(ImageReader):
+    def read_image(self, filename: str):
+        self.img = itk.imread(filename)
+
+    def get_meta_dict(self) -> List:
+        img_meta_dict = self.img.GetMetaDataDictionary()
+        meta_dict = dict()
+        for key in img_meta_dict.GetKeys():
+            # ignore deprecated, legacy members that cause issues
+            if key.startswith("ITK_original_"):
+                continue
+            meta_dict[key] = img_meta_dict[key]
+        meta_dict["origin"] = np.asarray(self.img.GetOrigin())
+        meta_dict["spacing"] = np.asarray(self.img.GetSpacing())
+        meta_dict["direction"] = itk.array_from_matrix(self.img.GetDirection())
+        return meta_dict
+
+    def get_affine(self) -> List:
+        """
+        Construct Affine matrix based on direction, spacing, origin information.
+        Refer to: https://github.com/RSIP-Vision/medio
+
+        """
+        direction = itk.array_from_matrix(self.img.GetDirection())
+        spacing = np.asarray(self.img.GetSpacing())
+        origin = np.asarray(self.img.GetOrigin())
+
+        direction = np.asarray(direction)
+        affine = np.eye(direction.shape[0] + 1)
+        affine[(slice(-1), slice(-1))] = direction @ np.diag(spacing)
+        affine[(slice(-1), -1)] = origin
+        return affine
+
+    def get_spatial_shape(self) -> List:
+        spatial_shape = list(itk.size(self.img))
+        spatial_shape.reverse()
+        return spatial_shape
+
+    def get_array_data(self) -> np.ndarray:
+        return itk.array_view_from_image(self.img)

monai/transforms/io/array.py

@@ -20,6 +20,7 @@
 from torch.utils.data._utils.collate import np_str_obj_array_pattern
 
 from monai.config import KeysCollection
+from monai.data.image_reader import ImageReader, ITKReader
 from monai.data.utils import correct_nifti_header_if_necessary
 from monai.transforms.compose import Transform
 from monai.utils import ensure_tuple, optional_import
@@ -28,6 +29,78 @@
 Image, _ = optional_import("PIL.Image")
 
 
+class LoadImage(Transform):
+    """
+    Load image file or files from provided path based on reader, default reader is ITK.
+    All the supported image formats of ITK:
+    https://github.com/InsightSoftwareConsortium/ITK/tree/master/Modules/IO
+    If loading a list of files, stack them together and add a new dimension as first dimension,
+    and use the meta data of the first image to represent the stacked result.
+
+    """
+
+    def __init__(
+        self,
+        reader: ImageReader = None,
+        suffixes: Optional[Union[str, Sequence[str]]] = None,
+        image_only: bool = False,
+        dtype: Optional[np.dtype] = np.float32,
+    ) -> None:
+        """
+        Args:
+            reader: use reader to load image file and meta data, default is ITK.
+            suffixes: file suffixes that supported by the reader, if None, support all kinds of files.
+            image_only: if True return only the image volume, otherwise return image data array and header dict.
+            dtype: if not None convert the loaded image to this data type.
+
+        Note:
+            The transform returns image data array if `image_only` is True,
+            or a tuple of two elements containing the data array, and the meta data in a dict format otherwise.
+
+        """
+        if reader is None:
+            reader = ITKReader()
+        self.reader = reader
+        self.suffixes = suffixes
+        self.image_only = image_only
+        self.dtype = dtype
+
+    def __call__(self, filename: Union[Sequence[Union[Path, str]], Path, str]):
+        """
+        Args:
+            filename: path file or file-like object or a list of files.
+        """
+        filename = ensure_tuple(filename)
+        img_array = list()
+        compatible_meta: Dict = None
+        for name in filename:
+            if not self.reader.verify_suffix(name.split(".")[-1]):
+                raise RuntimeError("unsupported file format.")
+            self.reader.read_image(name)
+            img_array.append(self.reader.get_array_data().astype(dtype=self.dtype))
+            if self.image_only:
+                continue
+
+            header = self.reader.get_meta_dict()
+            header["filename_or_obj"] = name
+            header["affine"] = self.reader.get_affine()
+            header["original_affine"] = header["affine"].copy()
+            header["spatial_shape"] = self.reader.get_spatial_shape()
+
+            if compatible_meta is None:
+                compatible_meta = header
+            else:
+                assert np.allclose(
+                    header["affine"], compatible_meta["affine"]
+                ), "affine data of all images should be same."
+
+        img_array = np.stack(img_array, axis=0) if len(img_array) > 1 else img_array[0]
+        self.reader.uncache()
+        if self.image_only:
+            return img_array
+        return img_array, compatible_meta
+
+
 class LoadNifti(Transform):
     """
     Load Nifti format file or files from provided path. If loading a list of