This is basically my loader. It just loads a .h5 file with the 1 channel for each one the 4 modalities and applies some data transformation.
class BraTS2018(Dataset):
""" BraTS2018 Dataset """
def __init__(self, base_dir=None, img_type=None, split='train', num=None, transform=None):
self._base_dir = base_dir
self.transform = transform
self.sample_list = []
self.split = split
self.img_type = img_type
print(self._base_dir)
train_path = self._base_dir+'/train.list'
val_path = self._base_dir+'/val.list'
test_path = self._base_dir+'/test.list'
if self.split == 'train':
with open(train_path, 'r') as f:
self.image_list = f.readlines()
elif self.split == 'val':
with open(val_path, 'r') as f:
self.image_list = f.readlines()
elif self.split == 'test':
with open(test_path, 'r') as f:
self.image_list = f.readlines()
self.image_list = [item.replace('\n', '').split(",")[0] for item in self.image_list]
if num is not None:
self.image_list = self.image_list[:num]
print("total {} samples".format(len(self.image_list)))
def __len__(self):
return len(self.image_list)
def __getitem__(self, idx):
case = self.image_list[idx]
h5f = h5py.File(self._base_dir + "/{}.h5".format(case), "r")
flair, t1, t2, ce, seg = h5f['image'][:], h5f['t1'][:], h5f['t2'][:] , h5f['ce'][:], h5f['label'][:]
input_image = np.zeros((4,flair.shape[0], flair.shape[1], flair.shape[2]))
input_image[0,:,:,:] = flair
input_image[1,:,:,:] = t1
input_image[2,:,:,:] = t2
input_image[3,:,:,:] = ce
seg[seg==4]=3
# image, label = flair[:], seg[:]
sample = {'image': input_image, 'label': seg.astype(np.uint8), 'name': case}
if self.transform:
sample = self.transform(sample)
return sample
class CenterCrop(object):
def __init__(self, output_size):
self.output_size = output_size
def __call__(self, sample):
image, label = sample['image'], sample['label']
# pad the sample if necessary
if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
self.output_size[2]:
pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)],
mode='constant', constant_values=0)
label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)],
mode='constant', constant_values=0)
(w, h, d) = image.shape
w1 = int(round((w - self.output_size[0]) / 2.))
h1 = int(round((h - self.output_size[1]) / 2.))
d1 = int(round((d - self.output_size[2]) / 2.))
label = label[w1:w1 + self.output_size[0], h1:h1 +
self.output_size[1], d1:d1 + self.output_size[2]]
image = image[w1:w1 + self.output_size[0], h1:h1 +
self.output_size[1], d1:d1 + self.output_size[2]]
return {'image': image, 'label': label}
class RandomCrop(object):
"""
Crop randomly the image in a sample
Args:
output_size (int): Desired output size
"""
def __init__(self, output_size, with_sdf=False):
self.output_size = output_size
self.with_sdf = with_sdf
def __call__(self, sample):
image, label = sample['image'], sample['label']
# pad the sample if necessary
if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
self.output_size[2]:
pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
new_array = np.zeros((4, label.shape[0]+2*pw, label.shape[1]+2*ph, label.shape[2]+2*pd))
new_array[0,:,:,:] = np.pad(image[0,:,:,:], [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
new_array[1,:,:,:] = np.pad(image[1,:,:,:], [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
new_array[2,:,:,:] = np.pad(image[2,:,:,:], [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
new_array[3,:,:,:] = np.pad(image[3,:,:,:], [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
image = new_array
(_, w, h, d) = image.shape
# if np.random.uniform() > 0.33:
# w1 = np.random.randint((w - self.output_size[0])//4, 3*(w - self.output_size[0])//4)
# h1 = np.random.randint((h - self.output_size[1])//4, 3*(h - self.output_size[1])//4)
# else:
w1 = np.random.randint(0, w - self.output_size[0])
h1 = np.random.randint(0, h - self.output_size[1])
d1 = np.random.randint(0, d - self.output_size[2])
label = label[w1:w1 + self.output_size[0], h1:h1 +
self.output_size[1], d1:d1 + self.output_size[2]]
# label[] = label[w1:w1 + self.output_size[0], h1:h1 +
# self.output_size[1], d1:d1 + self.output_size[2]]
image = image[:, w1:w1 + self.output_size[0], h1:h1 +
self.output_size[1], d1:d1 + self.output_size[2]]
return {'image': image, 'label': label}
class RandomRotFlip(object):
"""
Crop randomly flip the dataset in a sample
Args:
output_size (int): Desired output size
"""
def __call__(self, sample):
flair, t1, t2, ce, label = sample['image'][0,:,:,:], sample['image'][1,:,:,:], sample['image'][2,:,:,:], sample['image'][3,:,:,:], sample['label']
k = np.random.randint(0, 4)
flair = np.rot90(flair, k)
ce = np.rot90(ce, k)
t1 = np.rot90(t1, k)
t2 = np.rot90(t2, k)
label = np.rot90(label, k)
axis = np.random.randint(0, 2)
flair = np.flip(flair, axis=axis).copy()
ce = np.flip(ce, axis=axis).copy()
t1 = np.flip(t1, axis=axis).copy()
t2 = np.flip(t2, axis=axis).copy()
label = np.flip(label, axis=axis).copy()
input_image = np.zeros((4,flair.shape[0], flair.shape[1], flair.shape[2]))
input_image[0,:,:,:] = flair
input_image[1,:,:,:] = t1
input_image[2,:,:,:] = t2
input_image[3,:,:,:] = ce
return {'image': input_image, 'label': label}
class RandomNoise(object):
def __init__(self, mu=0, sigma=0.1):
self.mu = mu
self.sigma = sigma
def __call__(self, sample):
image, label = sample['image'], sample['label']
noise = np.clip(self.sigma * np.random.randn(
image.shape[0], image.shape[1], image.shape[2]), -2*self.sigma, 2*self.sigma)
noise = noise + self.mu
image = image + noise
return {'image': image, 'label': label}
class CreateOnehotLabel(object):
def __init__(self, num_classes):
self.num_classes = num_classes
def __call__(self, sample):
image, label = sample['image'], sample['label']
onehot_label = np.zeros(
(self.num_classes, label.shape[0], label.shape[1], label.shape[2]), dtype=np.float32)
for i in range(self.num_classes):
onehot_label[i, :, :, :] = (label == i).astype(np.float32)
return {'image': image, 'label': label, 'onehot_label': onehot_label}
class ToTensor(object):
"""Convert ndarrays in sample to Tensors."""
def __call__(self, sample):
image = sample['image']
return {'image': torch.from_numpy(image).float(), 'label': torch.from_numpy(sample['label']).long()}
class TwoStreamBatchSampler(Sampler):
"""Iterate two sets of indices
An 'epoch' is one iteration through the primary indices.
During the epoch, the secondary indices are iterated through
as many times as needed.
"""
def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size):
self.primary_indices = primary_indices
self.secondary_indices = secondary_indices
self.secondary_batch_size = secondary_batch_size
self.primary_batch_size = batch_size - secondary_batch_size
assert len(self.primary_indices) >= self.primary_batch_size > 0
assert len(self.secondary_indices) >= self.secondary_batch_size > 0
def __iter__(self):
primary_iter = iterate_once(self.primary_indices)
secondary_iter = iterate_eternally(self.secondary_indices)
return (
primary_batch + secondary_batch
for (primary_batch, secondary_batch)
in zip(grouper(primary_iter, self.primary_batch_size),
grouper(secondary_iter, self.secondary_batch_size))
)
def __len__(self):
return len(self.primary_indices) // self.primary_batch_size
def iterate_once(iterable):
return np.random.permutation(iterable)
def iterate_eternally(indices):
def infinite_shuffles():
while True:
yield np.random.permutation(indices)
return itertools.chain.from_iterable(infinite_shuffles())
def grouper(iterable, n):
"Collect data into fixed-length chunks or blocks"
# grouper('ABCDEFG', 3) --> ABC DEF"
args = [iter(iterable)] * n
return zip(*args)
Regarding the output shape and the label shape, the prediction of the model has a shape of (2,4,128,128,128), (B, C, H, W,D), respectively, and the label has a shape of (2,128,128,128), (B, H, W, D).
Is this a problem when calculating the Cross Entropy loss?