def generate_train_batch(self):
patient_ix = np.random.choice(range(len(self._data)), self.BATCH_SIZE, replace=True)
patient_ids = [self._data.keys()[ix] for ix in patient_ix]
data, seg = [], []
for b in range(self.BATCH_SIZE):
patient = self._data[patient_ids[b]]
shp = patient['data'].shape
if self.dim == 2:
slice_ix = np.random.choice(range(shp[2]))
tmp_data = resize_image_by_padding(patient['data'][:, :, slice_ix], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1])), pad_value=0)
tmp_seg = resize_image_by_padding(patient['seg'][:, :, slice_ix], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1])), pad_value=0)
data.append(center_crop_2D_image(tmp_data, self.pre_crop_size)[np.newaxis])
seg.append(center_crop_2D_image(tmp_seg, self.pre_crop_size)[np.newaxis])
elif self.dim == 3:
tmp_data = resize_image_by_padding(patient['data'], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1]),
max(shp[2], self.pre_crop_size[2])), pad_value=0)
tmp_seg = resize_image_by_padding(patient['seg'], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1]),
max(shp[2], self.pre_crop_size[2])), pad_value=0)
data.append(center_crop_3D_image(tmp_data, self.pre_crop_size)[np.newaxis])
seg.append(center_crop_3D_image(tmp_seg, self.pre_crop_size)[np.newaxis])
return {'data': np.array(data).astype('float32'), 'seg': np.array(seg).astype('float32'), 'pid': patient_ids}
def generate_train_batch(self):
if self.patient_ix == len(self.test_pids):
raise StopIteration
pid = self.test_pids[self.patient_ix]
patient = self._data[pid]
shp = patient['data'].shape
if self.dim == 2:
z_pre_crop_size = shp[2]
else:
z_pre_crop_size = max(shp[2], self.pre_crop_size[2])
data_arr = resize_image_by_padding(patient['data'], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1]), z_pre_crop_size), pad_value=0)
seg_arr = resize_image_by_padding(patient['seg'], (
max(shp[0], self.pre_crop_size[0]), max(shp[1], self.pre_crop_size[1]), z_pre_crop_size), pad_value=0)
data_arr = center_crop_3D_image(data_arr, (self.pre_crop_size[0], self.pre_crop_size[0], z_pre_crop_size))[np.newaxis]
seg_arr = center_crop_3D_image(seg_arr, (self.pre_crop_size[0], self.pre_crop_size[0], z_pre_crop_size))[np.newaxis]
if self.dim == 2:
data_arr = np.transpose(data_arr, axes=(3, 0, 1, 2))
seg_arr = np.transpose(seg_arr, axes=(3, 0, 1, 2))
else:
data_arr = data_arr[np.newaxis]
seg_arr = seg_arr[np.newaxis]
self.patient_ix += 1
return {'data': data_arr.astype('float32'), 'seg': seg_arr.astype('float32'), 'pid': pid}
def generate_train_batch(self):
data = np.zeros(
(self.BATCH_SIZE, 1, self.PATCH_SIZE[0], self.PATCH_SIZE[1]),
dtype=np.float32)
seg = np.zeros(
(self.BATCH_SIZE, 1, self.PATCH_SIZE[0], self.PATCH_SIZE[1]),
dtype=np.float32)
types = np.random.choice(['ed', 'es'], self.BATCH_SIZE,
True) # randomly choose es or ed
# print(self._data.keys())
patients = np.random.choice(
list(self._data.keys()), self.BATCH_SIZE,
True) # randomly choose 'batchsize' patients
pathologies = []
for nb in range(self.BATCH_SIZE):
"""
data[i]['ed_data'] = a[0, :]
data[i]['ed_gt'] = a[1, :]
data[i]['es_data'] = a[2, :]
data[i]['es_gt'] = a[3, :]
"""
shp = self._data[patients[nb]][types[nb] + '_data'].shape
slice_id = np.random.choice(
shp[0]
) # randomly choose one slice from total slice, here is 10
tmp_data = resize_image_by_padding(
self._data[patients[nb]][types[nb] + '_data'][slice_id],
(max(shp[1],
self.PATCH_SIZE[0]), max(shp[2], self.PATCH_SIZE[1])),
pad_value=0)
tmp_seg = resize_image_by_padding(
self._data[patients[nb]][types[nb] + '_gt'][slice_id],
(max(shp[1],
self.PATCH_SIZE[0]), max(shp[2], self.PATCH_SIZE[1])),
pad_value=0)
# not the most efficient way but whatever...
tmp = np.zeros((1, 2, tmp_data.shape[0], tmp_data.shape[1]))
tmp[0, 0] = tmp_data
tmp[0, 1] = tmp_seg
tmp = random_crop_2D_image_batched(tmp, self.PATCH_SIZE)
data[nb, 0] = tmp[0, 0]
seg[nb, 0] = tmp[0, 1]
pathologies.append(self._data[patients[nb]]['pathology'])
# print(data)
return {
'data': data,
'seg': seg,
'types': types,
'patient_ids': patients,
'pathologies': pathologies
}
def generate_train_batch(self):
data = np.zeros((self.BATCH_SIZE, 1, self.PATCH_SIZE[0],
self.PATCH_SIZE[1], self.PATCH_SIZE[2]),
dtype=np.float32)
seg = np.zeros((self.BATCH_SIZE, 1, self.PATCH_SIZE[0],
self.PATCH_SIZE[1], self.PATCH_SIZE[2]),
dtype=np.float32)
types = np.random.choice(['ed', 'es'], self.BATCH_SIZE, True)
patients = np.random.choice(self._data.keys(), self.BATCH_SIZE, True)
pathologies = []
for nb in range(self.BATCH_SIZE):
if np.any(
np.array(self._data[patients[nb]][types[nb] + '_data'].
shape) < np.array(self.PATCH_SIZE)):
shp = self._data[patients[nb]][types[nb] + '_data'].shape
tmp_data = resize_image_by_padding(
self._data[patients[nb]][types[nb] + '_data'],
(max(shp[0],
self.PATCH_SIZE[0]), max(shp[1], self.PATCH_SIZE[1]),
max(shp[2], self.PATCH_SIZE[2])),
pad_value=0)
tmp_seg = resize_image_by_padding(
self._data[patients[nb]][types[nb] + '_gt'],
(max(shp[0],
self.PATCH_SIZE[0]), max(shp[1], self.PATCH_SIZE[1]),
max(shp[2], self.PATCH_SIZE[2])),
pad_value=0)
else:
tmp_data = self._data[patients[nb]][types[nb] + '_data']
tmp_seg = self._data[patients[nb]][types[nb] + '_gt']
# not the most efficient way but whatever...
tmp = np.zeros((1, 2, tmp_data.shape[0], tmp_data.shape[1],
tmp_data.shape[2]))
tmp[0, 0] = tmp_data
tmp[0, 1] = tmp_seg
if self._random_crop:
tmp = random_crop_3D_image_batched(tmp, self.PATCH_SIZE)
else:
tmp = center_crop_3D_image_batched(tmp, self.PATCH_SIZE)
data[nb, 0] = tmp[0, 0]
seg[nb, 0] = tmp[0, 1]
pathologies.append(self._data[patients[nb]]['pathology'])
return {
'data': data,
'seg': seg,
'types': types,
'patient_ids': patients,
'pathologies': pathologies
}
def generate_train_batch(self):
# DataLoader has its own methods for selecting what patients to use next, see its Documentation
idx = self.get_indices()
patients_for_batch = [self._data[i] for i in idx]
print(idx)
patch_size = None
if self.patch_size == None:
shapes = []
for i in patients_for_batch:
patient_data, _ = self.load_patient(i)
shapes.append(patient_data[0].shape)
patch_size = np.max(shapes, 0)
patch_size = (patch_size[0] + 16 - patch_size[0] % 16,
patch_size[1] + 16 - patch_size[1] % 16,
patch_size[2] + 16 - patch_size[2] % 16)
else:
patch_size = self.patch_size
# initialize empty array for data and seg
data = np.zeros((self.batch_size, self.num_modalities, *patch_size),
dtype=np.float32)
seg = np.zeros((self.batch_size, 1, *patch_size), dtype=np.float32)
metadata = []
patient_names = []
# iterate over patients_for_batch and include them in the batch
for i, j in enumerate(patients_for_batch):
patient_data, patient_metadata = self.load_patient(j)
if self.patch_size is not None:
# this will only pad patient_data if its shape is smaller than self.patch_size
patient_data = pad_nd_image(patient_data, patch_size)
# now random crop to self.patch_size
# crop expects the data to be (b, c, x, y, z) but patient_data is (c, x, y, z) so we need to add one
# dummy dimension in order for it to work (@Todo, could be improved)
patient_data, patient_seg = crop(patient_data[:-1][None],
patient_data[-1:][None],
patch_size,
crop_type="random")
data[i] = patient_data[0]
seg[i] = patient_seg[0]
else:
for j in range(len(patient_data)):
if j != len(patient_data) - 1:
data[i][j] = resize_image_by_padding(
patient_data[j], patch_size)
else:
seg[i][0] = resize_image_by_padding(
patient_data[j], patch_size)
metadata.append(patient_metadata)
patient_names.append(j)
return {
'data': data,
'seg': seg,
'metadata': metadata,
'names': patient_names
}