"""Functions for reading MSISBI2015 NRRD data."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import glob
import math
import os.path
import pickle
import matplotlib.pyplot
from imageio import imwrite
from scipy.ndimage import zoom
from six.moves import xrange
from skimage.measure import label, regionprops
from utils.NII import *
from utils.image_utils import crop, crop_center
from utils.tfrecord_utils import *
matplotlib.pyplot.ion()
class MSISBI2015(object):
PROTOCOL_MAPPINGS = {'FLAIR': ['flair'], 'MPRAGE': ['mprage'], 'PD': ['pd'], 'T2': ['t2']}
SET_TYPES = ['TRAIN', 'VAL', 'TEST']
class Options(object):
def __init__(self):
self.dir = os.path.dirname(os.path.realpath(__file__))
self.numSamples = -1
self.partition = {'TRAIN': 0.7, 'VAL': 0.2, 'TEST': 0.1}
self.useCrops = False
self.cropType = 'random' # random or center
self.numRandomCropsPerSlice = 5
self.onlyPatchesWithLesions = False
self.rotations = 0
self.cropWidth = 128
self.cropHeight = 128
self.cache = False
self.sliceResolution = None # format: HxW
self.addInstanceNoise = False # Affects only the batch sampling. If True, a tiny bit of noise will be added to every batch
self.filterProtocol = None # FLAIR, T1, T2
self.filterType = "train" # train or test
self.axis = 'axial' # saggital, coronal or axial
self.debug = False
self.normalizationMethod = 'standardization'
self.sliceStart = 0
self.sliceEnd = 155
self.format = "raw" # raw or aligned; If aligned, nii-files will be crawled and loaded
self.skullStripping = True
self.viewMapping = {'saggital': 2, 'coronal': 1, 'axial': 0}
def __init__(self, options=Options()):
self.options = options
if options.cache and os.path.isfile(self.pckl_name()):
f = open(self.pckl_name(), 'rb')
tmp = pickle.load(f)
f.close()
self._epochs_completed = tmp._epochs_completed
self._index_in_epoch = tmp._index_in_epoch
self.patientsSplit = tmp.patients_split
self.patients = tmp.patients
self._images, self._labels, self._sets = read_tf_record(self.tfrecord_name())
self._epochs_completed = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}
self._index_in_epoch = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}
else:
# Collect all patients
self.patients = self._get_patients()
self.patientsSplit = {}
if not os.path.isfile(self.split_name()):
_numPatients = len(self.patients)
_ridx = numpy.random.permutation(_numPatients)
_already_taken = 0
for split in self.options.partition.keys():
if self.options.partition[split] <= 1.0:
num_patients_for_current_split = math.floor(self.options.partition[split] * _numPatients)
else:
num_patients_for_current_split = self.options.partition[split]
if num_patients_for_current_split > (_numPatients - _already_taken):
num_patients_for_current_split = _numPatients - _already_taken
self.patientsSplit[split] = _ridx[_already_taken:_already_taken + num_patients_for_current_split]
_already_taken += num_patients_for_current_split
f = open(self.split_name(), 'wb')
pickle.dump(self.patientsSplit, f)
f.close()
else:
f = open(self.split_name(), 'rb')
self.patientsSplit = pickle.load(f)
f.close()
self._create_numpy_arrays()
self._epochs_completed = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}
self._index_in_epoch = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}
if self.options.cache:
write_tf_record(self._images, self._labels, self._sets, self.tfrecord_name())
tmp = copy.copy(self)
tmp._images = None
tmp._labels = None
tmp._sets = None
f = open(self.pckl_name(), 'wb')
pickle.dump(tmp, f)
f.close()
def _create_numpy_arrays(self):
# Iterate over all patients and extract slices
_images = []
_labels = []
_sets = []
for p, patient in enumerate(self.patients):
if p in self.patientsSplit['TRAIN']:
_set_of_current_patient = MSISBI2015.SET_TYPES.index('TRAIN')
elif p in self.patientsSplit['VAL']:
_set_of_current_patient = MSISBI2015.SET_TYPES.index('VAL')
elif p in self.patientsSplit['TEST']:
_set_of_current_patient = MSISBI2015.SET_TYPES.index('TEST')
for n, nii_filename in enumerate(patient['filtered_files']):
# try:
_images_tmp, _labels_tmp = self.gather_data(patient, nii_filename)
_images += _images_tmp
_labels += _labels_tmp
_sets += [_set_of_current_patient] * len(_images_tmp)
self._images = numpy.array(_images).astype(numpy.float32)
self._labels = numpy.array(_labels).astype(numpy.float32)
if self._images.ndim < 4:
self._images = numpy.expand_dims(self._images, 3)
self._sets = numpy.array(_sets).astype(numpy.int32)
def gather_data(self, patient, nii_filename):
_images = []
_labels = []
nii, nii_seg, nii_skullmap = self.load_volume_and_groundtruth(nii_filename, patient)
# Iterate over all slices and collect them
# We only want to select in the range from 15 to 125 (in axial view)
for s in xrange(self.options.sliceStart, min(self.options.sliceEnd, nii.num_slices_along_axis(self.options.axis))):
if 0 < self.options.numSamples < len(_images):
break
slice_data = nii.get_slice(s, self.options.axis)
slice_seg = nii_seg.get_slice(s, self.options.axis)
# Skip the slice if it is "empty"
if numpy.percentile(slice_data, 90) < 0.2:
continue
if self.options.sliceResolution is not None:
# Pad withzeros to top and bottom, if the image is too small
if slice_data.shape[0] < self.options.sliceResolution[0]:
before_y = math.floor((self.options.sliceResolution[0] - slice_data.shape[0]) / 2.0)
after_y = math.ceil((self.options.sliceResolution[0] - slice_data.shape[0]) / 2.0)
if slice_data.shape[1] < self.options.sliceResolution[1]:
before_x = math.floor((self.options.sliceResolution[1] - slice_data.shape[1]) / 2.0)
after_x = math.ceil((self.options.sliceResolution[1] - slice_data.shape[1]) / 2.0)
if slice_data.shape[0] < self.options.sliceResolution[0] or slice_data.shape[1] < self.options.sliceResolution[1]:
slice_data = np.pad(slice_data, ((before_y, after_y), (before_x, after_x)), 'constant', constant_values=(0, 0))
slice_seg = np.pad(slice_seg, ((before_y, after_y), (before_x, after_x)), 'constant', constant_values=(0, 0))
slice_data = zoom(slice_data, float(self.options.sliceResolution[0]) / float(
slice_data.shape[0]))
slice_seg = zoom(
slice_seg, float(self.options.sliceResolution[0]) / float(slice_seg.shape[0]), mode="nearest"
)
slice_seg[slice_seg < 0.9] = 0.0
slice_seg[slice_seg >= 0.9] = 1.0
# assert numpy.max(slice_data) <= 1.0, "Resized slice range is outside [0; 1]!"
# Either collect crops
if self.options.useCrops:
if self.options.cropType == 'random':
rx = numpy.random.randint(0, high=(slice_data.shape[1] - self.options.cropWidth),
size=self.options.numRandomCropsPerSlice)
ry = numpy.random.randint(0, high=(slice_data.shape[0] - self.options.cropHeight),
size=self.options.numRandomCropsPerSlice)
for r in range(self.options.numRandomCropsPerSlice):
_images.append(crop(slice_data, ry(r), rx(r), self.options.cropHeight, self.options.cropWidth))
_labels.append(crop(slice_data, ry(r), rx(r), self.options.cropHeight, self.options.cropWidth))
elif self.options.cropType == 'center':
slice_data_cropped = crop_center(slice_data, self.options.cropWidth, self.options.cropHeight)
slice_seg_cropped = crop_center(slice_seg, self.options.cropWidth, self.options.cropHeight)
_images.append(slice_data_cropped)
_labels.append(slice_seg_cropped)
elif self.options.cropType == 'lesions':
cc_slice = label(slice_seg)
props = regionprops(cc_slice)
if len(props) > 0:
for prop in props:
cx = prop['centroid'][1]
cy = prop['centroid'][0]
if cy < self.options.cropHeight // 2:
cy = self.options.cropHeight // 2
if cy > (slice_data.shape[0] - (self.options.cropHeight // 2)):
cy = (slice_data.shape[0] - (self.options.cropHeight // 2))
if cx < self.options.cropWidth // 2:
cx = self.options.cropWidth // 2
if cx > (slice_data.shape[1] - (self.options.cropWidth // 2)):
cx = (slice_data.shape[1] - (self.options.cropWidth // 2))
image_crop = crop(slice_data, int(cy) - (self.options.cropHeight // 2), int(cx) - (self.options.cropWidth // 2),
self.options.cropHeight, self.options.cropWidth)
seg_crop = crop(slice_seg, int(cy) - (self.options.cropHeight // 2), int(cx) - (self.options.cropWidth // 2),
self.options.cropHeight, self.options.cropWidth)
if image_crop.shape[0] != self.options.cropHeight or image_crop.shape[1] != self.options.cropWidth:
continue
_images.append(image_crop)
_labels.append(seg_crop)
# Or whole slices
else:
_images.append(slice_data)
_labels.append(slice_seg)
return _images, _labels
def load_volume_and_groundtruth(self, nii_filename, patient):
# Load the nrrd
try:
nii = NII(nii_filename)
nii_groundtruth = NII(patient['groundtruth'])
nii.denoise()
nii.set_view_mapping(self.options.viewMapping)
except:
print('MSISBI2015: Failed to open file ' + nii_filename)
# Make sure ground-truth is binary and nrrd doesnt have NaNs
nii.data[np.isnan(nii.data)] = 0.0
nii_groundtruth.data[nii_groundtruth.data < 0.9] = 0.0
nii_groundtruth.data[nii_groundtruth.data >= 0.9] = 1.0
# Do skull-stripping, if desired
if self.options.skullStripping:
try:
nii_skullmap = NII(patient['skullmap'])
nii_skullmap.set_view_mapping(self.options.viewMapping)
nii.apply_skullmap(nii_skullmap)
except:
print('MSISBI2015: Failed to open file ' + patient['skullmap'] + ', skipping skullremoval')
# In-place normalize the loaded volume
nii.normalize(method=self.options.normalizationMethod, lowerpercentile=0, upperpercentile=99.8)
# nii_skullmap.data = nii_skullmap.data > 0.0
return nii, nii_groundtruth, nii_skullmap
# Hidden helper function, not supposed to be called from outside!
def _get_patients(self):
return MSISBI2015.get_patients(self.options)
@staticmethod
def get_patients(options):
folders = ["training01", "training02", "training03", "training04", "training05"]
# Iterate over all folders in folders and collect all patients
patients = []
for f, folder in enumerate(folders):
# Get all files that can be used for training and validation
_patients = glob.glob(os.path.join(options.dir, folder, "preprocessed", folder + "_*_flair_pp.nii"))
for p, pname in enumerate(_patients):
patient = {}
_tmp = os.path.normpath(pname).split(os.path.sep)
patient['name'] = _tmp[-1].replace("_flair_pp.nii", "")
patient['fullpath'] = os.path.join(options.dir, folder, "preprocessed")
patient["filtered_files"] = []
for protocol, protocol_array in MSISBI2015.PROTOCOL_MAPPINGS.items():
if len(options.filterProtocols) > 0 and protocol not in options.filterProtocols:
continue
else:
if options.format == "raw":
patient[protocol] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '_pp.nii')
elif options.format == "aligned":
patient[protocol] = os.path.join(options.dir, folder, "preprocessed",
patient['name'] + '_' + protocol_array[0] + '.aligned.nii.gz')
if len(options.filterProtocols) > 0 and protocol not in options.filterProtocols:
continue
else:
if options.format == "raw":
patient["filtered_files"] += [
os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '_pp.nii')]
elif options.format == "aligned":
patient["filtered_files"] += [
os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '.aligned.nii.gz')]
if options.format == "raw":
patient['groundtruth'] = os.path.join(options.dir, folder, "masks", patient['name'] + "_mask1.nii")
patient['skullmap'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_skullmap.nii.gz")
elif options.format == "aligned":
patient['groundtruth'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_mask1.aligned.nii.gz")
patient['skullmap'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_skullmap.aligned.nii.gz")
# Append to the list of all patients
patients.append(patient)
return patients
# Returns the indices of patients which belong to either TRAIN, VAL or TEST. Your choice
def get_patient_idx(self, split='TRAIN'):
return self.patientsSplit[split]
def get_patient_split(self):
return self.patientsSplit
@property
def images(self):
return self._images
def get_images(self, set=None):
_setIdx = MSISBI2015.SET_TYPES.index(set)
images_in_set = numpy.where(self._sets == _setIdx)[0]
return self._images[images_in_set]
def get_image(self, i):
return self._images[i, :, :, :]
def get_label(self, i):
return self._labels[i, :, :, :]
def get_patient(self, i):
return self.patients[i]
@property
def labels(self):
return self._labels
@property
def sets(self):
return self._sets
@property
def meta(self):
return self._meta
@property
def num_examples(self):
return self._images.shape[0]
@property
def width(self):
return self._images.shape[2]
@property
def height(self):
return self._images.shape[1]
@property
def num_channels(self):
return self._images.shape[3]
@property
def epochs_completed(self):
return self._epochs_completed
def name(self):
_name = "MSISBI2015"
if self.options.numSamples > 0:
_name += '_n{}'.format(self.options.numSamples)
_name += "_p{}-{}".format(self.options.partition['TRAIN'], self.options.partition['VAL'])
if self.options.useCrops:
_name += "_{}crops{}x{}".format(self.options.cropType, self.options.cropWidth, self.options.cropHeight)
if self.options.cropType == "random":
_name += "_{}cropsPerSlice".format(self.options.numRandomCropsPerSlice)
if self.options.sliceResolution is not None:
_name += "_res{}x{}".format(self.options.sliceResolution[0], self.options.sliceResolution[1])
_name += "_{}".format(self.options.format)
return _name
def split_name(self):
return os.path.join(self.dir(), 'split-{}-{}.pckl'.format(self.options.partition['TRAIN'], self.options.partition['VAL']))
def pckl_name(self):
return os.path.join(self.dir(), self.name() + ".pckl")
def tfrecord_name(self):
return os.path.join(self.dir(), self.name() + ".tfrecord")
def dir(self):
return self.options.dir
def export_slices(self, dir):
for i in range(self.num_examples):
imwrite(os.path.join(dir, '{}.png'.format(i)), np.squeeze(self.get_image(i) * 255).astype('uint8'))
def visualize(self, pause=1):
f, (ax1, ax2) = matplotlib.pyplot.subplots(1, 2)
images_tmp, labels_tmp, _ = self.next_batch(10)
for i in range(images_tmp.shape[0]):
img = numpy.squeeze(images_tmp[i])
lbl = numpy.squeeze(labels_tmp[i])
ax1.imshow(img)
ax1.set_title('Patch')
ax2.imshow(lbl)
ax2.set_title('Groundtruth')
matplotlib.pyplot.pause(pause)
def num_batches(self, batchsize, set='TRAIN'):
_setIdx = MSISBI2015.SET_TYPES.index(set)
images_in_set = numpy.where(self._sets == _setIdx)[0]
return len(images_in_set) // batchsize
def next_batch(self, batch_size, shuffle=True, set='TRAIN', return_brainmask=True):
"""Return the next `batch_size` examples from this data set."""
_setIdx = MSISBI2015.SET_TYPES.index(set)
images_in_set = numpy.where(self._sets == _setIdx)[0]
samples_in_set = len(images_in_set)
start = self._index_in_epoch[set]
# Shuffle for the first epoch
if self._epochs_completed[set] == 0 and start == 0 and shuffle:
perm0 = numpy.arange(samples_in_set)
numpy.random.shuffle(perm0)
self._images[images_in_set] = self.images[images_in_set[perm0]]
self._labels[images_in_set] = self.labels[images_in_set[perm0]]
self._sets[images_in_set] = self.sets[images_in_set[perm0]]
# Go to the next epoch
if start + batch_size > samples_in_set:
# Finished epoch
self._epochs_completed[set] += 1
# Get the rest examples in this epoch
rest_num_examples = samples_in_set - start
images_rest_part = self._images[images_in_set[start:samples_in_set]]
labels_rest_part = self._labels[images_in_set[start:samples_in_set]]
# Shuffle the data
if shuffle:
perm = numpy.arange(samples_in_set)
numpy.random.shuffle(perm)
self._images[images_in_set] = self.images[images_in_set[perm]]
self._labels[images_in_set] = self.labels[images_in_set[perm]]
self._sets[images_in_set] = self.sets[images_in_set[perm]]
# Start next epoch
start = 0
self._index_in_epoch[set] = batch_size - rest_num_examples
end = self._index_in_epoch[set]
images_new_part = self._images[images_in_set[start:end]]
labels_new_part = self._labels[images_in_set[start:end]]
images_tmp = numpy.concatenate((images_rest_part, images_new_part), axis=0)
labels_tmp = numpy.concatenate((labels_rest_part, labels_new_part), axis=0)
else:
self._index_in_epoch[set] += batch_size
end = self._index_in_epoch[set]
images_tmp = self._images[images_in_set[start:end]]
labels_tmp = self._labels[images_in_set[start:end]]
if self.options.addInstanceNoise:
noise = numpy.random.normal(0, 0.01, images_tmp.shape)
images_tmp += noise
# Check the batch
assert images_tmp.size, "The batch is empty!"
assert labels_tmp.size, "The labels of the current batch are empty!"
if return_brainmask:
brainmasks = images_tmp > 0.05
else:
brainmasks = None
return images_tmp, labels_tmp, brainmasks
