def extract_crop(self, i, label=None, vol_size=(64, 64, 64)):
def get_output(oname):
name_suffix = oname + str(vol_size[0])
output_image = self.get_imagetype_path(name_suffix)
if not os.path.exists(output_image):
os.makedirs(output_image)
return output_image, name_suffix
name = self.get_image_name(i)
print(name)
output_image, image_name_suffix = get_output(
configs.get('images_crop_prefix'))
output_labelmap, labelmap_name_suffix = get_output(
configs.get('labelmaps_crop_prefix'))
lmap = self.load_labelmap(i)
if label is not None:
labels = [label]
else:
labels = list(np.unique(lmap).flatten())
labels.remove(0)
sitk_image = self.load_sitk_image(i)
spacing = sitk_image.GetSpacing()
img = self.get_array_from_sitk_image(sitk_image)
for l in labels:
image, labelmap = extract_vol_at_label(img,
self.load_labelmap(i),
label=l,
vol_size=vol_size)
labelmap = (labelmap == l).astype(np.uint8)
image = sitk.GetImageFromArray(image)
labelmap = sitk.GetImageFromArray(labelmap)
image.SetSpacing(spacing)
labelmap.SetSpacing(spacing)
suffix = '_' + str(l)
sitk.WriteImage(
image,
os.path.join(
output_image,
name + suffix + '_' + image_name_suffix + '.nii.gz'))
sitk.WriteImage(
labelmap,
os.path.join(
output_labelmap,
name + suffix + '_' + labelmap_name_suffix + '.nii.gz'))
def load_labelmap_crop(self, img_idx, vol_size=(64, 64, 64), label=1):
name = self.get_image_name(img_idx)
name_suffix = configs.get('labelmaps_crop_prefix') + str(vol_size[0])
output = self.get_imagetype_path(name_suffix)
path = os.path.join(
output, name + '_' + str(label) + '_' + name_suffix + '.nii.gz')
return self._load_image_from_disk(path)
def setup(self):
if self.dir_path is None:
return
dataset_path = Path(self.dir_path)
image_type_paths = {
image_type_path.name: image_type_path
for image_type_path in dataset_path.iterdir()
}
image_types = list(image_type_paths.keys())
image_types.remove('images')
image_types = ['images'] + image_types
for image_type in image_types:
image_type_path = image_type_paths[image_type]
image_type = image_type_path.name
self.image_type_dirs.add(image_type)
image_type = configs.get('remap_dirs',
{}).get(image_type, image_type)
for image_path in (image_type_path /
self.get_spacing_dirname()).glob('*' +
self.ext):
name = image_path.name.replace(self.ext, '')
for existing_name in self.dataframe.index:
if existing_name in name: # check if subset of existing name
name = existing_name
self.dataframe.loc[name,
f'{image_type}_path'] = str(image_path)
if self.images_only:
self.dataframe = self.dataframe[['image_path']].dropna()
else:
self.dataframe.dropna(inplace=True)
def load2d_slices(self, label):
s = self.spacing
try:
s = s[0]
except:
pass
path = os.path.join(
self.dir_path, self.name + '_label' + str(label) + '_spacing' +
str(s) + '_2dslices.npz')
if os.path.exists(path):
slices = np.load(path)
else:
print('{} does not exist. Extracting...'.format(path))
self.export_2d_slices(self.dir_path, label)
slices = np.load(path)
return slices[configs.get('images_dir')], slices[configs.get(
'labelmaps_dir')]
def get_spacing_dirname(self, spacing=None):
if spacing is None:
spacing = self.spacing
if type(spacing) in [int, float]:
if isinstance(spacing, float) and spacing.is_integer():
spacing = int(spacing)
spacing = [spacing]
if sum(spacing) <= 0:
spacing_dirname = configs.get('native_images_dir')
elif len(spacing) == 1:
spacing_dirname = configs.get(
'subsampled_images_dir_prefix') + str(spacing[0]) + 'mm'
else:
spacing_str = ''
for s in spacing:
spacing_str += str(s) + '-'
spacing_str = spacing_str[:-1]
spacing_dirname = configs.get(
'subsampled_images_dir_prefix') + spacing_str + 'mm'
return spacing_dirname
def get_root_path(self):
return configs.get('root_path')
def load_dataset(name, **kwargs):
dataset = get_dataset_info(name)
dataset['dir_path'] = os.path.join(configs.get('root_path'), dataset['subpath'])
dataset.update(kwargs)
return MIReader.from_dict(**dataset)