def transform_and_copy(data, cahce_dir):
copy_dir = os.path.join(cahce_dir, 'copied_images')
if not os.path.exists(copy_dir):
os.mkdir(copy_dir)
copy_list_path = os.path.join(copy_dir, 'copied_images.npy')
if not os.path.exists(copy_list_path):
print("transforming and copying images...")
imageLoader = LoadImage()
to_copy_list = [x for x in data if int(x['_label']) == 1]
mul = 1 #int(len(data)/len(to_copy_list) - 1)
rand_x_flip = RandFlip(spatial_axis=0, prob=0.50)
rand_y_flip = RandFlip(spatial_axis=1, prob=0.50)
rand_z_flip = RandFlip(spatial_axis=2, prob=0.50)
rand_affine = RandAffine(prob=1.0,
rotate_range=(0, 0, np.pi / 10),
shear_range=(0.12, 0.12, 0.0),
translate_range=(0, 0, 0),
scale_range=(0.12, 0.12, 0.0),
padding_mode="zeros")
rand_gaussian_noise = RandGaussianNoise(prob=0.5, mean=0.0, std=0.05)
transform = Compose([
AddChannel(),
rand_x_flip,
rand_y_flip,
rand_z_flip,
rand_affine,
SqueezeDim(),
])
copy_list = []
n = len(to_copy_list)
for i in range(len(to_copy_list)):
print(f'Copying image {i+1}/{n}', end="\r")
to_copy = to_copy_list[i]
image_file = to_copy['image']
_image_file = replace_suffix(image_file, '.nii.gz', '')
label = to_copy['label']
_label = to_copy['_label']
image_data, _ = imageLoader(image_file)
seg_file = to_copy['seg']
seg_data, _ = nrrd.read(seg_file)
for i in range(mul):
rand_seed = np.random.randint(1e8)
transform.set_random_state(seed=rand_seed)
new_image_data = rand_gaussian_noise(
np.array(transform(image_data)))
transform.set_random_state(seed=rand_seed)
new_seg_data = np.array(transform(seg_data))
#multi_slice_viewer(image_data, image_file)
#multi_slice_viewer(seg_data, seg_file)
#seg_image = MaskIntensity(seg_data)(image_data)
#multi_slice_viewer(seg_image, seg_file)
image_basename = os.path.basename(_image_file)
seg_basename = image_basename + f'_seg_{i}.nrrd'
image_basename = image_basename + f'_{i}.nii.gz'
new_image_file = os.path.join(copy_dir, image_basename)
write_nifti(new_image_data, new_image_file, resample=False)
new_seg_file = os.path.join(copy_dir, seg_basename)
nrrd.write(new_seg_file, new_seg_data)
copy_list.append({
'image': new_image_file,
'seg': new_seg_file,
'label': label,
'_label': _label
})
np.save(copy_list_path, copy_list)
print("done transforming and copying!")
copy_list = np.load(copy_list_path, allow_pickle=True)
return copy_list
def large_image_splitter(data, cache_dir, num_splits, only_label_one=False):
print("Splitting large images...")
len_old = len(data)
print("original data len:", len_old)
split_images_dir = os.path.join(cache_dir, 'split_images')
split_images = os.path.join(split_images_dir, 'split_images.npy')
def _replace_in_data(split_images, num_splits):
new_images = []
for image in data:
new_images.append(image)
for s in split_images:
source_image = s['source']
if image['_label'] == 0 and only_label_one is True:
break
if image['image'] == source_image:
#new_images.pop()
for i in range(min(num_splits, len(s["splits"]))):
new_images.append(s["splits"][i])
break
return new_images
if os.path.exists(split_images):
new_images = np.load(split_images, allow_pickle=True)
"""for s in new_images:
print("split image:", s["source"], end='\r')"""
out_data = _replace_in_data(new_images, num_splits)
else:
if not os.path.exists(split_images_dir):
os.mkdir(split_images_dir)
new_images = []
imageLoader = LoadImage()
for image in data:
image_data, _ = imageLoader(image["image"])
seg_data, _ = nrrd.read(image['seg'])
label = image['_label']
z_len = image_data.shape[2]
if z_len > 200:
count = z_len // 80
print("splitting image:",
image["image"],
f"into {count} parts",
"shape:",
image_data.shape,
end='\r')
split_image_list = [
image_data[:, :, idz::count] for idz in range(count)
]
split_seg_list = [
seg_data[:, :, idz::count] for idz in range(count)
]
new_image = {'source': image["image"], 'splits': []}
for i in range(count):
image_file = os.path.basename(
replace_suffix(image["image"], '.nii.gz', ''))
image_file = os.path.join(split_images_dir,
image_file + f'_{i}.nii.gz')
seg_file = os.path.basename(
replace_suffix(image["seg"], '.nrrd', ''))
seg_file = os.path.join(split_images_dir,
seg_file + f'_seg_{i}.nrrd')
split_image = np.array(split_image_list[i])
split_seg = np.array(split_seg_list[i], dtype=np.uint8)
rand_affine = RandAffine(prob=1.0,
rotate_range=(0, 0, np.pi / 16),
shear_range=(0.07, 0.07, 0.0),
translate_range=(0, 0, 0),
scale_range=(0.07, 0.07, 0.0),
padding_mode="zeros")
transform = Compose([
AddChannel(),
rand_affine,
SqueezeDim(),
])
rand_seed = np.random.randint(1e8)
transform.set_random_state(seed=rand_seed)
split_image = transform(split_image).detach().cpu().numpy()
transform.set_random_state(seed=rand_seed)
split_seg = transform(split_seg).detach().cpu().numpy()
write_nifti(split_image, image_file, resample=False)
nrrd.write(seg_file, split_seg)
new_image['splits'].append({
'image': image_file,
'label': image['label'],
'_label': image['_label'],
'seg': seg_file,
'w': False
})
new_images.append(new_image)
np.save(split_images, new_images)
out_data = _replace_in_data(new_images, num_splits)
print("new data len:", len(out_data))
return out_data