simpleitk.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156
  1. # -*- coding: utf-8 -*-
  2. # imageio is distributed under the terms of the (new) BSD License.
  3. """Read/Write images using SimpleITK.
  4. Backend: `Insight Toolkit <https://itk.org/>`_
  5. .. note::
  6. To use this plugin you have to install its backend::
  7. pip install imageio[itk]
  8. The ItkFormat uses the ITK or SimpleITK library to support a range of
  9. ITK-related formats. It also supports a few common formats (e.g. PNG and JPEG).
  10. Parameters
  11. ----------
  12. None
  13. """
  14. from ..core import Format, has_module
  15. _itk = None # Defer loading to load_lib() function.
  16. def load_lib():
  17. global _itk, _read_function, _write_function
  18. try:
  19. import itk as _itk
  20. _read_function = _itk.imread
  21. _write_function = _itk.imwrite
  22. except ImportError:
  23. try:
  24. import SimpleITK as _itk
  25. _read_function = _itk.ReadImage
  26. _write_function = _itk.WriteImage
  27. except ImportError:
  28. raise ImportError(
  29. "itk could not be found. "
  30. "Please try "
  31. " python -m pip install itk "
  32. "or "
  33. " python -m pip install simpleitk "
  34. "or refer to "
  35. " https://itkpythonpackage.readthedocs.io/ "
  36. "for further instructions."
  37. )
  38. return _itk
  39. # Split up in real ITK and all supported formats.
  40. ITK_FORMATS = (
  41. ".gipl",
  42. ".ipl",
  43. ".mha",
  44. ".mhd",
  45. ".nhdr",
  46. "nia",
  47. "hdr",
  48. ".nrrd",
  49. ".nii",
  50. ".nii.gz",
  51. ".img",
  52. ".img.gz",
  53. ".vtk",
  54. "hdf5",
  55. "lsm",
  56. "mnc",
  57. "mnc2",
  58. "mgh",
  59. "mnc",
  60. "pic",
  61. )
  62. ALL_FORMATS = ITK_FORMATS + (
  63. ".bmp",
  64. ".jpeg",
  65. ".jpg",
  66. ".png",
  67. ".tiff",
  68. ".tif",
  69. ".dicom",
  70. ".dcm",
  71. ".gdcm",
  72. )
  73. class ItkFormat(Format):
  74. """See :mod:`imageio.plugins.simpleitk`"""
  75. def _can_read(self, request):
  76. # If the request is a format that only this plugin can handle,
  77. # we report that we can do it; a useful error will be raised
  78. # when simpleitk is not installed. For the more common formats
  79. # we only report that we can read if the library is installed.
  80. if request.extension in ITK_FORMATS:
  81. return True
  82. if has_module("itk.ImageIOBase") or has_module("SimpleITK"):
  83. return request.extension in ALL_FORMATS
  84. def _can_write(self, request):
  85. if request.extension in ITK_FORMATS:
  86. return True
  87. if has_module("itk.ImageIOBase") or has_module("SimpleITK"):
  88. return request.extension in ALL_FORMATS
  89. # -- reader
  90. class Reader(Format.Reader):
  91. def _open(self, pixel_type=None, fallback_only=None, **kwargs):
  92. if not _itk:
  93. load_lib()
  94. args = ()
  95. if pixel_type is not None:
  96. args += (pixel_type,)
  97. if fallback_only is not None:
  98. args += (fallback_only,)
  99. self._img = _read_function(self.request.get_local_filename(), *args)
  100. def _get_length(self):
  101. return 1
  102. def _close(self):
  103. pass
  104. def _get_data(self, index):
  105. # Get data
  106. if index != 0:
  107. error_msg = "Index out of range while reading from itk file"
  108. raise IndexError(error_msg)
  109. # Return array and empty meta data
  110. return _itk.GetArrayFromImage(self._img), {}
  111. def _get_meta_data(self, index):
  112. error_msg = "The itk plugin does not support meta data, currently."
  113. raise RuntimeError(error_msg)
  114. # -- writer
  115. class Writer(Format.Writer):
  116. def _open(self):
  117. if not _itk:
  118. load_lib()
  119. def _close(self):
  120. pass
  121. def _append_data(self, im, meta):
  122. _itk_img = _itk.GetImageFromArray(im)
  123. _write_function(_itk_img, self.request.get_local_filename())
  124. def set_meta_data(self, meta):
  125. error_msg = "The itk plugin does not support meta data, currently."
  126. raise RuntimeError(error_msg)