forked from 170010011/fr
163 lines
4.2 KiB
Python
163 lines
4.2 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
# imageio is distributed under the terms of the (new) BSD License.
|
||
|
|
||
|
""" Storage of image data in multiple formats.
|
||
|
"""
|
||
|
|
||
|
from .. import formats
|
||
|
from ..core import Format, has_module
|
||
|
|
||
|
_itk = None # Defer loading to load_lib() function.
|
||
|
|
||
|
|
||
|
def load_lib():
|
||
|
global _itk, _read_function, _write_function
|
||
|
try:
|
||
|
import itk as _itk
|
||
|
|
||
|
_read_function = _itk.imread
|
||
|
_write_function = _itk.imwrite
|
||
|
except ImportError:
|
||
|
try:
|
||
|
import SimpleITK as _itk
|
||
|
|
||
|
_read_function = _itk.ReadImage
|
||
|
_write_function = _itk.WriteImage
|
||
|
except ImportError:
|
||
|
raise ImportError(
|
||
|
"itk could not be found. "
|
||
|
"Please try "
|
||
|
" python -m pip install itk "
|
||
|
"or "
|
||
|
" python -m pip install simpleitk "
|
||
|
"or refer to "
|
||
|
" https://itkpythonpackage.readthedocs.io/ "
|
||
|
"for further instructions."
|
||
|
)
|
||
|
return _itk
|
||
|
|
||
|
|
||
|
# Split up in real ITK and all supported formats.
|
||
|
ITK_FORMATS = (
|
||
|
".gipl",
|
||
|
".ipl",
|
||
|
".mha",
|
||
|
".mhd",
|
||
|
".nhdr",
|
||
|
"nia",
|
||
|
"hdr",
|
||
|
".nrrd",
|
||
|
".nii",
|
||
|
".nii.gz",
|
||
|
".img",
|
||
|
".img.gz",
|
||
|
".vtk",
|
||
|
"hdf5",
|
||
|
"lsm",
|
||
|
"mnc",
|
||
|
"mnc2",
|
||
|
"mgh",
|
||
|
"mnc",
|
||
|
"pic",
|
||
|
)
|
||
|
ALL_FORMATS = ITK_FORMATS + (
|
||
|
".bmp",
|
||
|
".jpeg",
|
||
|
".jpg",
|
||
|
".png",
|
||
|
".tiff",
|
||
|
".tif",
|
||
|
".dicom",
|
||
|
".dcm",
|
||
|
".gdcm",
|
||
|
)
|
||
|
|
||
|
|
||
|
class ItkFormat(Format):
|
||
|
""" The ItkFormat uses the ITK or SimpleITK library to support a range of
|
||
|
ITK-related formats. It also supports a few common formats that are
|
||
|
also supported by the freeimage plugin (e.g. PNG and JPEG).
|
||
|
|
||
|
This format requires the ``itk`` or ``SimpleITK`` package.
|
||
|
|
||
|
Parameters for reading
|
||
|
----------------------
|
||
|
None.
|
||
|
|
||
|
Parameters for saving
|
||
|
---------------------
|
||
|
None.
|
||
|
|
||
|
"""
|
||
|
|
||
|
def _can_read(self, request):
|
||
|
# If the request is a format that only this plugin can handle,
|
||
|
# we report that we can do it; a useful error will be raised
|
||
|
# when simpleitk is not installed. For the more common formats
|
||
|
# we only report that we can read if the library is installed.
|
||
|
if request.extension in ITK_FORMATS:
|
||
|
return True
|
||
|
if has_module("itk.ImageIOBase") or has_module("SimpleITK"):
|
||
|
return request.extension in ALL_FORMATS
|
||
|
|
||
|
def _can_write(self, request):
|
||
|
if request.extension in ITK_FORMATS:
|
||
|
return True
|
||
|
if has_module("itk.ImageIOBase") or has_module("SimpleITK"):
|
||
|
return request.extension in ALL_FORMATS
|
||
|
|
||
|
# -- reader
|
||
|
|
||
|
class Reader(Format.Reader):
|
||
|
def _open(self, pixel_type=None, fallback_only=None, **kwargs):
|
||
|
if not _itk:
|
||
|
load_lib()
|
||
|
args = ()
|
||
|
if pixel_type is not None:
|
||
|
args += (pixel_type,)
|
||
|
if fallback_only is not None:
|
||
|
args += (fallback_only,)
|
||
|
self._img = _read_function(self.request.get_local_filename(), *args)
|
||
|
|
||
|
def _get_length(self):
|
||
|
return 1
|
||
|
|
||
|
def _close(self):
|
||
|
pass
|
||
|
|
||
|
def _get_data(self, index):
|
||
|
# Get data
|
||
|
if index != 0:
|
||
|
error_msg = "Index out of range while reading from itk file"
|
||
|
raise IndexError(error_msg)
|
||
|
|
||
|
# Return array and empty meta data
|
||
|
return _itk.GetArrayFromImage(self._img), {}
|
||
|
|
||
|
def _get_meta_data(self, index):
|
||
|
error_msg = "The itk plugin does not support meta data, currently."
|
||
|
raise RuntimeError(error_msg)
|
||
|
|
||
|
# -- writer
|
||
|
class Writer(Format.Writer):
|
||
|
def _open(self):
|
||
|
if not _itk:
|
||
|
load_lib()
|
||
|
|
||
|
def _close(self):
|
||
|
pass
|
||
|
|
||
|
def _append_data(self, im, meta):
|
||
|
_itk_img = _itk.GetImageFromArray(im)
|
||
|
_write_function(_itk_img, self.request.get_local_filename())
|
||
|
|
||
|
def set_meta_data(self, meta):
|
||
|
error_msg = "The itk plugin does not support meta data, currently."
|
||
|
raise RuntimeError(error_msg)
|
||
|
|
||
|
|
||
|
# Register
|
||
|
title = "Insight Segmentation and Registration Toolkit (ITK) format"
|
||
|
format = ItkFormat("itk", title, " ".join(ALL_FORMATS), "iIvV")
|
||
|
formats.add_format(format)
|