"""Functions for reading BRAINWEB NII 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 cv2
import matplotlib.pyplot
from imageio import imwrite
from scipy.ndimage import rotate
from utils.MINC import *
from utils.image_utils import crop, crop_center
from utils.tfrecord_utils import *
class BRAINWEB(object):
FILTER_TYPES = ['NORMAL', 'MILDMS', 'MODERATEMS', 'SEVEREMS']
SET_TYPES = ['TRAIN', 'VAL', 'TEST']
LABELS = {'BACKGROUND': 0, 'CSF': 1, 'GM': 2, 'WM': 3, 'FAT': 4, 'MUSCLE': 5, 'SKIN': 6, 'SKULL': 7, 'GLIALMATTER': 8, 'CONNECTIVE': 9, 'LESION': 10}
VIEW_MAPPING = {'saggital': 0, 'coronal': 1, 'axial': 2}
PROTOCOL_MAPPINGS = {'FLAIR': 'flair*', 'T2': 't2*'}
class Options(object):
def __init__(self):
self.description = None
self.dir = os.path.dirname(os.path.realpath(__file__))
self.folderNormal = 'normal'
self.folderMildMS = os.path.join('lesions', 'mild')
self.folderModerateMS = os.path.join('lesions', 'moderate')
self.folderSevereMS = os.path.join('lesions', 'severe')
self.folderGT = 'groundtruth'
self.numSamples = -1
self.partition = {'TRAIN': 0.6, 'VAL': 0.15, 'TEST': 0.25}
self.sliceStart = 20
self.sliceEnd = 140
self.useCrops = False
self.cropType = 'random' # random or center
self.numRandomCropsPerSlice = 5
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 # T2 or FLAIR only, not implemented for now
self.filterType = None # MILDMS, MODERATEMS, SEVEREMS, NORMAL
self.axis = 'axial' # saggital, coronal or axial
self.debug = False
self.normalizationMethod = 'standardization'
self.skullRemoval = False
self.backgroundRemoval = False
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.patients = self._get_patients()
self._images, self._labels, self._sets = read_tf_record(self.tfrecord_name())
f = open(self.split_name(), 'rb')
self.patients_split = pickle.load(f)
f.close()
if not os.path.exists(self.split_name() + ".deprecated"):
os.rename(self.split_name(), self.split_name() + ".deprecated")
self._convert_patient_split()
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.patients_split = {} # Here we will later store the info whether a patient belongs to train, val or test
# Determine Train, Val & Test set based on patients
if not os.path.isfile(self.split_name()):
_num_patients = len(self.patients)
_ridx = numpy.random.permutation(_num_patients)
_already_taken = 0
for split in self.options.partition.keys():
if 1.0 >= self.options.partition[split] > 0.0:
num_patients_for_current_split = max(1, math.floor(self.options.partition[split] * _num_patients))
else:
num_patients_for_current_split = int(self.options.partition[split])
if num_patients_for_current_split > (_num_patients - _already_taken):
num_patients_for_current_split = _num_patients - _already_taken
self.patients_split[split] = _ridx[_already_taken:_already_taken + num_patients_for_current_split]
_already_taken += num_patients_for_current_split
self._convert_patient_split() # NEW! We have a new format for storing hte patientsSplit which is OS agnostic.
else:
f = open(self.split_name(), 'rb')
self.patients_split = pickle.load(f)
f.close()
self._convert_patient_split() # NEW! We have a new format for storing hte patientsSplit which is OS agnostic.
# Iterate over all patients and the filtered NII files and extract slices
_images = []
_labels = []
_sets = []
for p, patient in enumerate(self.patients):
if patient["name"] in self.patients_split['TRAIN']:
_set_of_current_patient = BRAINWEB.SET_TYPES.index('TRAIN')
elif patient["name"] in self.patients_split['VAL']:
_set_of_current_patient = BRAINWEB.SET_TYPES.index('VAL')
elif patient["name"] in self.patients_split['TEST']:
_set_of_current_patient = BRAINWEB.SET_TYPES.index('TEST')
minc, minc_seg, minc_skullmap = self.load_volume_and_groundtruth(patient["filtered_files"][0], patient)
# Iterate over all slices and collect them
for s in range(self.options.sliceStart, min(self.options.sliceEnd, minc.num_slices_along_axis(self.options.axis))):
if 0 < self.options.numSamples < len(_images):
break
slice_data = minc.get_slice(s, self.options.axis)
slice_seg = minc_seg.get_slice(s, self.options.axis)
# Skip the slice if it is entirely black
if numpy.unique(slice_data).size == 1:
continue
# assert numpy.max(slice_data) <= 1.0, "Slice range is outside [0; 1]!"
if self.options.sliceResolution is not None:
# If the images are too big in resolution, do downsampling
if slice_data.shape[0] > self.options.sliceResolution[0] or slice_data.shape[1] > self.options.sliceResolution[1]:
slice_data = cv2.resize(slice_data, tuple(self.options.sliceResolution))
slice_seg = cv2.resize(slice_seg, tuple(self.options.sliceResolution), interpolation=cv2.INTER_NEAREST)
# Otherwise, do zero padding
else:
tmp_slice = numpy.zeros(self.options.sliceResolution)
tmp_slice_seg = numpy.zeros(self.options.sliceResolution)
start_x = (self.options.sliceResolution[1] - slice_data.shape[1]) // 2
start_y = (self.options.sliceResolution[0] - slice_data.shape[0]) // 2
end_x = start_x + slice_data.shape[1]
end_y = start_y + slice_data.shape[0]
tmp_slice[start_y:end_y, start_x:end_x] = slice_data
tmp_slice_seg[start_y:end_y, start_x:end_x] = slice_seg
slice_data = tmp_slice
slice_seg = tmp_slice_seg
for angle in self.options.rotations:
if angle != 0:
slice_data_rotated = rotate(slice_data, angle, reshape=False)
slice_seg_rotated = rotate(slice_seg, angle, reshape=False, mode='nearest')
else:
slice_data_rotated = slice_data
slice_seg_rotated = slice_seg
# Either collect crops
if self.options.useCrops:
if self.options.cropType == 'random':
rx = numpy.random.randint(0, high=(slice_data_rotated.shape[1] - self.options.cropWidth),
size=self.options.numRandomCropsPerSlice)
ry = numpy.random.randint(0, high=(slice_data_rotated.shape[0] - self.options.cropHeight),
size=self.options.numRandomCropsPerSlice)
for r in range(self.options.numRandomCropsPerSlice):
_images.append(crop(slice_data_rotated, ry[r], rx[r], self.options.cropHeight, self.options.cropWidth))
_labels.append(crop(slice_data_rotated, ry[r], rx[r], self.options.cropHeight, self.options.cropWidth))
_sets.append(_set_of_current_patient)
elif self.options.cropType == 'center':
slice_data_cropped = crop_center(slice_data_rotated, self.options.cropWidth, self.options.cropHeight)
slice_seg_cropped = crop_center(slice_seg_rotated, self.options.cropWidth, self.options.cropHeight)
_images.append(slice_data_cropped)
_labels.append(slice_seg_cropped)
_sets.append(_set_of_current_patient)
# Or whole slices
else:
_images.append(slice_data_rotated)
_labels.append(slice_seg_rotated)
_sets.append(_set_of_current_patient)
self._images = numpy.array(_images).astype(numpy.float32)
self._labels = numpy.array(_labels).astype(numpy.float32)
# assert numpy.max(self._images) <= 1.0, "MINC range is outside [0; 1]!"
if self._images.ndim < 4:
self._images = numpy.expand_dims(self._images, 3)
self._sets = numpy.array(_sets).astype(numpy.int32)
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 _get_patients(self):
return BRAINWEB.get_patients(self.options)
@staticmethod
def get_patients(options):
minc_folders = [options.folderNormal, options.folderMildMS, options.folderModerateMS, options.folderSevereMS]
# Iterate over all folders and collect patients
patients = []
for n, minc_folder in enumerate(minc_folders):
if minc_folder == options.folderNormal:
_type = 'NORMAL'
elif minc_folder == options.folderMildMS:
_type = 'MILDMS'
elif minc_folder == options.folderModerateMS:
_type = 'MODERATEMS'
elif minc_folder == options.folderSevereMS:
_type = 'SEVEREMS'
# Continue with the next patient if the current one is not part of the desired types
if _type not in options.filterType:
continue
if options.filterProtocol:
_regex = BRAINWEB.PROTOCOL_MAPPINGS[options.filterProtocol] + ".mnc.gz"
else:
_regex = "*.mnc.gz"
_files = glob.glob(os.path.join(options.dir, minc_folder, _regex))
for f, fname in enumerate(_files):
patient = {
'name': os.path.basename(fname),
'type': _type,
'fullpath': fname
}
patient['filtered_files'] = patient['fullpath']
if patient['type'] == 'NORMAL':
patient['groundtruth_filename'] = os.path.join(options.dir, options.folderGT, 'normal.mnc.gz')
elif patient['type'] == 'MILDMS':
patient['groundtruth_filename'] = os.path.join(options.dir, options.folderGT, 'mild_lesions.mnc.gz')
elif patient['type'] == 'MODERATEMS':
patient['groundtruth_filename'] = os.path.join(options.dir, options.folderGT, 'moderate_lesions.mnc.gz')
elif patient['type'] == 'SEVEREMS':
patient['groundtruth_filename'] = os.path.join(options.dir, options.folderGT, 'severe_lesions.mnc.gz')
patients.append(patient)
return patients
def load_volume_and_groundtruth(self, minc_filename, patient):
minc_filename = patient['fullpath']
try:
minc = MINC(minc_filename) # NII also works with MINC
minc.set_view_mapping(BRAINWEB.VIEW_MAPPING)
except:
print('BRAINWEB: Failed to open file ' + minc_filename)
# Try to load the segmentation ground-truth
minc_seg_path = patient["groundtruth_filename"]
minc_seg = MINC(minc_seg_path)
skullmap = MINC(minc_seg_path)
skullmap.data = (skullmap.data * 0.0) + 1.0
skullmap.set_view_mapping(BRAINWEB.VIEW_MAPPING)
minc_seg.set_view_mapping(BRAINWEB.VIEW_MAPPING)
# If desired, compute the skullmap
if self.options.skullRemoval:
skullmap.data[minc_seg.data == BRAINWEB.LABELS['FAT']] = 0
skullmap.data[minc_seg.data == BRAINWEB.LABELS['MUSCLE']] = 0
skullmap.data[minc_seg.data == BRAINWEB.LABELS['SKIN']] = 0
skullmap.data[minc_seg.data == BRAINWEB.LABELS['SKULL']] = 0
skullmap.data[minc_seg.data == BRAINWEB.LABELS['CONNECTIVE']] = 0
if self.options.backgroundRemoval:
skullmap.data[minc_seg.data == BRAINWEB.LABELS['BACKGROUND']] = 0
# Binarize minc_seg
lesion_idx = (minc_seg.data == BRAINWEB.LABELS['LESION'])
nonlesion_idx = (minc_seg.data != BRAINWEB.LABELS['LESION'])
minc_seg.data[lesion_idx] = 1
minc_seg.data[nonlesion_idx] = 0
if self.options.skullRemoval or self.options.backgroundRemoval:
minc.apply_skullmap(skullmap)
# In-place normalize the loaded volume
minc.normalize(method=self.options.normalizationMethod, lowerpercentile=0.0, upperpercentile=99.8)
# 99.8 percentile described in LG Ny´ul, Jayaram K Udupa, and Xuan Zhang.
# New variants of a method of MRI scale standardization.
# IEEE transactions on medical imaging, 19(2):143–150, 2000.
# assert numpy.max(minc.getData()) <= 1.0, "MINC range is outside [0; 1]!"
return minc, minc_seg, skullmap
# Returns the indices of patients which belong to either TRAIN, VAL or TEST. Your choice
def get_patient_idx(self, split='TRAIN'):
idx = []
for pidx, patient in enumerate(self.patients):
if patient["name"] in self.patients_split[split]:
idx += [pidx]
return idx
def get_patient_split(self):
return self.patients_split
@property
def images(self):
return self._images
def get_images(self, set=None):
_setIdx = BRAINWEB.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, :, :, :]
@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 = "BRAINWEB"
if self.options.description:
_name += "_{}".format(self.options.description)
if self.options.numSamples > 0:
_name += '_n{}'.format(self.options.numSamples)
_name += "_p{}-{}-{}".format(self.options.partition['TRAIN'], self.options.partition['VAL'], self.options.partition['TEST'])
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])
if self.options.skullRemoval:
_name += "_noSkull"
if self.options.backgroundRemoval:
_name += "_noBackground"
return _name
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 split_name(self):
return os.path.join(self.dir(),
'split-{}-{}-{}.pckl'.format(self.options.partition['TRAIN'], self.options.partition['VAL'], self.options.partition['TEST']))
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, set='TRAIN'):
f, (ax1, ax2) = matplotlib.pyplot.subplots(1, 2)
images_tmp, labels_tmp, _ = self.next_batch(10, set=set)
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 = BRAINWEB.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=False):
"""Return the next `batch_size` examples from this data set."""
_setIdx = BRAINWEB.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 == 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 = np.copy(labels_tmp)
brainmasks[brainmasks == BRAINWEB.LABELS['FAT']] = 0
brainmasks[brainmasks == BRAINWEB.LABELS['MUSCLE']] = 0
brainmasks[brainmasks == BRAINWEB.LABELS['SKIN']] = 0
brainmasks[brainmasks == BRAINWEB.LABELS['SKULL']] = 0
brainmasks[brainmasks == BRAINWEB.LABELS['CONNECTIVE']] = 0
brainmasks[brainmasks == BRAINWEB.LABELS['BACKGROUND']] = 0
brainmasks[brainmasks > 0] = 1
return images_tmp, labels_tmp, brainmasks
return images_tmp, labels_tmp, None
def _convert_patient_split(self):
for split in self.patients_split.keys():
_list_of_patient_names = []
for pidx in self.patients_split[split]:
if not isinstance(pidx, str):
_list_of_patient_names += [self.patients[pidx]['name']]
else:
_list_of_patient_names = self.patients_split[split]
break
self.patients_split[split] = _list_of_patient_names
f = open(self.split_name(), 'wb')
pickle.dump(self.patients_split, f)
f.close()
