def __getitem__(self, i):
patient, slice_idx = self._find_image_from_index(i)
if self.current_patient != patient:
self.image, self.segmentation = self._load_image(patient)
self.current_patient = patient
start = 0
middle_slice = self.n_slices // 2
im_slice = crop_to_bbox(self.image,
(0, slice_idx - middle_slice, start, start, 1,
self.n_slices, 512, 512))
seg_slice = crop_to_bbox(self.segmentation,
(0, slice_idx, start, start, 3, 1, 512, 512))
sample = {"image": im_slice.squeeze(0), "target": seg_slice.squeeze(1)}
sample = self.transform(sample)
im_slice = np.expand_dims(sample["image"], 0)
seg_slice = np.expand_dims(sample["target"], 1)
assert (
0 not in seg_slice.shape
), f"Segmentation slice has dimension of size 0: {seg_slice.shape}"
return im_slice, seg_slice
def __getitem__(self, i):
patient, slice_idx = self._find_image_from_index(i)
if self.current_patient != patient:
self.cbct, self.ct, self.segmentation_cbct, self.segmentation_ct = self._load_images(
patient)
self.current_patient = patient
logger.debug(
f"GETITEM Image and segmentation shapes:\nCBCT: {self.cbct.shape}\nCBCT seg: {self.segmentation_cbct.shape}"
)
start_ct = int((self.ct.shape[2] - 512) / 2)
start_cbct = int((self.cbct.shape[2] - 512) / 2)
middle_slice = self.n_slices // 2
ct_slice = crop_to_bbox(self.ct,
(0, slice_idx - middle_slice, start_ct,
start_ct, 1, self.n_slices, 512, 512))
ct_seg_slice = crop_to_bbox(self.segmentation_ct,
(0, slice_idx, 0, 0, 3, 1, 512, 512))
cbct_slice = crop_to_bbox(self.cbct,
(0, slice_idx - middle_slice, start_cbct,
start_cbct, 1, self.n_slices, 512, 512))
cbct_seg_slice = crop_to_bbox(self.segmentation_cbct,
(0, slice_idx, 0, 0, 3, 1, 512, 512))
ct_sample = {
"image": ct_slice.squeeze(0),
"target": ct_seg_slice.squeeze(1)
}
ct_sample = self.transform(ct_sample)
ct_slice = np.expand_dims(ct_sample["image"], 0)
ct_seg_slice = np.expand_dims(ct_sample["target"], 1)
cbct_sample = {
"image": cbct_slice.squeeze(0),
"target": cbct_seg_slice.squeeze(1)
}
cbct_sample = self.transform(cbct_sample)
cbct_slice = np.expand_dims(cbct_sample["image"], 0)
cbct_seg_slice = np.expand_dims(cbct_sample["target"], 1)
assert (
0 not in ct_seg_slice.shape and 0 not in cbct_seg_slice.shape
), f"Segmentation slice has dimension of size 0: {ct_seg_slice.shape}"
if self.cbct_only:
return cbct_slice, cbct_seg_slice
if self.return_meta:
return self.meta, patient, cbct_slice, ct_slice, cbct_seg_slice, ct_seg_slice
if self.return_patient:
return patient, cbct_slice, ct_slice, cbct_seg_slice, ct_seg_slice
return cbct_slice, ct_slice, cbct_seg_slice, ct_seg_slice
def _get_segmentation_ct(self, segmentations):
if len(segmentations) == 2:
seg_bladder = read_image(segmentations[0], no_meta=True)
seg_cervix_uterus = read_image(segmentations[1], no_meta=True)
start = int((seg_bladder.shape[2] - 512) / 2)
seg_bladder = crop_to_bbox(
seg_bladder, (0, start, start, seg_bladder.shape[0], 512, 512))
seg_cervix_uterus = crop_to_bbox(
seg_cervix_uterus,
(0, start, start, seg_cervix_uterus.shape[0], 512, 512))
all_segs = seg_bladder + seg_cervix_uterus
other = all_segs < 1
segs = [seg_bladder, seg_cervix_uterus, other]
segmentation = np.stack(segs).astype(int)
return segmentation
def __call__(self, sample):
image = sample['image']
target = sample['target']
if self.plane not in ['inplane']:
raise NotImplementedError()
image = np.swapaxes(image, 0, 2)
target = np.swapaxes(target, 0, 2)
image_resized = np.swapaxes(
rescale(image,
self.zoom_factor,
order=1,
mode='constant',
clip=False,
multichannel=True,
anti_aliasing=False), 2, 0)
target_resized = np.swapaxes(
rescale(target,
self.zoom_factor,
order=0,
mode='constant',
clip=False,
multichannel=True,
preserve_range=True,
anti_aliasing=False), 2, 0)
final_image_size = np.asarray(image.shape)
final_target_size = np.asarray(target.shape)
image_crop_offset = (np.asarray(image_resized.shape) -
final_image_size) // 2
target_crop_offset = (np.asarray(target_resized.shape) -
final_target_size) // 2
image = crop_to_bbox(image_resized,
combine_bbox(image_crop_offset, final_image_size))
target = crop_to_bbox(
target_resized, combine_bbox(target_crop_offset,
target_crop_offset))
sample['image'] = image
sample['target'] = target
return sample
def _get_segmentation(self, segmentations):
if len(segmentations) == 2:
seg_bladder = read_image(segmentations[0], no_meta=True)
seg_cervix_uterus = read_image(segmentations[1], no_meta=True)
all_segs = seg_bladder + seg_cervix_uterus
# Combine cervix and uterus segmentation
elif len(segmentations) == 3:
seg_bladder = read_image(segmentations[0], no_meta=True)
seg_cervix = read_image(segmentations[1], no_meta=True)
seg_uterus = read_image(segmentations[2], no_meta=True)
seg_cervix_uterus = (seg_cervix | seg_uterus)
all_segs = seg_bladder + seg_cervix + seg_uterus
start = int((all_segs.shape[1] - 512) / 2)
seg_bladder = crop_to_bbox(
seg_bladder, (0, start, start, seg_bladder.shape[0], 512, 512))
seg_cervix_uterus = crop_to_bbox(
seg_cervix_uterus,
(0, start, start, seg_cervix_uterus.shape[0], 512, 512))
all_segs = crop_to_bbox(all_segs,
(0, start, start, all_segs.shape[0], 512, 512))
other = all_segs < 1
segs = [seg_bladder, seg_cervix_uterus, other]
segmentation = np.stack(segs).astype(int)
return segmentation