"""NIfTI I/O.
This module contains NIfTI input/output helpers.
"""
import os
import nibabel as nib
import numpy as np
from dosma.data_io.format_io import DataReader, DataWriter, ImageDataFormat
from dosma.data_io.med_volume import MedicalVolume
from dosma.defaults import AFFINE_DECIMAL_PRECISION, SCANNER_ORIGIN_DECIMAL_PRECISION
from dosma.utils import io_utils
__all__ = ["NiftiReader", "NiftiWriter"]
[docs]class NiftiReader(DataReader):
"""A class for reading NIfTI files.
"""
data_format_code = ImageDataFormat.nifti
def __normalize_affine(self, affine):
# determine vector for through-plane pixel direction (k)
# 1. Normalize k_vector by magnitude
# 2. Multiply by magnitude given by SpacingBetweenSlices field
# These actions are done to avoid rounding errors that might result from float subtraction
aff = np.array(affine)
i_vec = np.round(np.array(aff[:3, 0]), AFFINE_DECIMAL_PRECISION)
j_vec = np.round(np.array(aff[:3, 1]), AFFINE_DECIMAL_PRECISION)
k_vec = np.round(np.array(aff[:3, 2]), AFFINE_DECIMAL_PRECISION)
aff[:3, 0] = i_vec
aff[:3, 1] = j_vec
aff[:3, 2] = k_vec
aff[:3, 3] = np.round(aff[:3, 3], SCANNER_ORIGIN_DECIMAL_PRECISION)
return aff
[docs] def load(self, file_path):
"""Load volume from NIfTI file path.
A NIfTI file should only correspond to one volume.
Args:
file_path (str): File path to NIfTI file.
Returns:
MedicalVolume: Loaded volume.
Raises:
FileNotFoundError: If `file_path` not found.
ValueError: If `file_path` does not end in a supported NIfTI extension.
"""
if not os.path.isfile(file_path):
raise FileNotFoundError("{} not found".format(file_path))
if not self.data_format_code.is_filetype(file_path):
raise ValueError(
"{} must be a file with extension '.nii' or '.nii.gz'".format(file_path)
)
nib_img = nib.load(file_path)
nib_img_affine = nib_img.affine
nib_img_affine = self.__normalize_affine(nib_img_affine)
np_img = nib_img.get_fdata()
return MedicalVolume(np_img, nib_img_affine)
[docs]class NiftiWriter(DataWriter):
"""A class for writing volumes in NIfTI format.
"""
data_format_code = ImageDataFormat.nifti
[docs] def save(self, volume: MedicalVolume, file_path: str):
"""Save volume in NIfTI format,
Args:
volume (MedicalVolume): Volume to save.
file_path (str): File path to NIfTI file.
Raises:
ValueError: If `file_path` does not end in a supported NIfTI extension.
"""
if not self.data_format_code.is_filetype(file_path):
raise ValueError(
"{} must be a file with extension '.nii' or '.nii.gz'".format(file_path)
)
# Create dir if does not exist
io_utils.mkdirs(os.path.dirname(file_path))
nib_affine = volume.affine
np_im = volume.volume
nib_img = nib.Nifti1Image(np_im, nib_affine)
nib.save(nib_img, file_path)