def restore(self, predicted_array_list):
predicted_array = np.zeros(self.label.GetSize()[::-1])
z_size = self.label.GetSize()[2] - self.label_patch_width
indices = [i for i in range(0, z_size + 1, self.z_slide)]
with tqdm(total=len(indices), desc="Restoring image...",
ncols=60) as pbar:
for pre_array, index in zip(predicted_array_list, indices):
z_slice = slice(index, index + self.label_patch_width)
predicted_array[z_slice, ...] = pre_array
pbar.update(1)
predicted = getImageWithMeta(predicted_array, self.label)
predicted = cropping(predicted, self.lower_pad_size[1].tolist(),
self.upper_pad_size[1].tolist())
if (self.diff > 0).any():
lower_crop_size = (self.diff // 2).tolist()
upper_crop_size = ((self.diff + 1) // 2).tolist()
predicted = cropping(predicted, lower_crop_size, upper_crop_size)
else:
lower_pad_size = (abs(self.diff) // 2).tolist()
upper_pad_size = ((abs(self.diff) + 1) // 2).tolist()
predicted = padding(predicted, lower_pad_size, upper_pad_size)
return predicted
def restore(self, predictArrayList):
labelArray = sitk.GetArrayFromImage(self.label)
outputArray = np.zeros_like(labelArray)
for i in range(2):
length = len(predictArrayList[i])
largestSlice = self.meta[i]["largestSlice"]
axialSlice = self.meta[i]["axialSlice"]
paddingSize = self.paddingSize
largestArray = outputArray[largestSlice, ...]
p = (paddingSize, paddingSize)
largestArray = np.pad(largestArray, [p, p, (0, 0)], "minimum")
largestArray = largestArray[..., axialSlice]
for x in tqdm(range(length), desc="Restoring images...", ncols=60):
pre = getImageWithMeta(predictArrayList[i][x], self.sliceImage)
sagittalSlice = self.meta[i]["sagittalSlice"][x]
coronalSlice = self.meta[i]["coronalSlice"][x]
size = self.meta[i]["size"][x]
pre = resampleSize(pre, size[::-1], is_label=True)
preArray = sitk.GetArrayFromImage(pre)
largestArray[sagittalSlice, coronalSlice, x] = preArray
largestArray = largestArray[paddingSize:-paddingSize,
paddingSize:-paddingSize, :]
if not self.noFlip:
if i == 1:
largestArray = largestArray[::-1, ...]
outputArray[largestSlice, :, axialSlice] += largestArray
return outputArray
def restore(self, predict_array_list):
predict_array = np.zeros_like(self.label_array)
size = np.array(self.label_array.shape) - self.label_array_patch_size
indices = []
for i in range(self.label_array.ndim):
r = range(0, size[i] + 1, self.slide[i])
indices.append(r)
indices = np.array([i for i in product(*indices)])
with tqdm(total=len(predict_array_list),
desc="Restoring image...",
ncols=60) as pbar:
for pre_array, idx in zip(predict_array_list,
indices[self.masked_indices]):
slices = []
for i in range(self.label_array.ndim):
s = slice(idx[i], idx[i] + self.label_array_patch_size[i])
slices.append(s)
slices = tuple(slices)
predict_array[slices] = pre_array
pbar.update(1)
predict_array = croppingForNumpy(predict_array,
self.lower_pad_size[1].tolist(),
self.upper_pad_size[1].tolist())
predict = getImageWithMeta(predict_array, self.org)
predict.SetOrigin(self.org.GetOrigin())
return predict
def restore(self, predict_array_list):
predicted_array = np.zeros(self.image.GetSize()[::-1])
""" Make patch. """
total = self.z_length // self.slide[2]
with tqdm(total=total, desc="Restoring images...", ncols=60) as pbar:
for z, predict_array in zip(range(0, self.z_length, self.slide[2]),
predict_array_list):
z_slice = slice(z, z + self.patch_size[2])
predicted_array[z_slice, ...] = predict_array
pbar.update(1)
predicted = getImageWithMeta(predicted_array, self.image)
predicted = cropping(predicted, self.axial_lower_pad_size[0].tolist(),
self.axial_upper_pad_size[0].tolist())
""" Raw data is clipped or padded for required_shape. """
if (self.diff > 0).any():
lower_crop_size = (abs(self.diff) // 2).tolist()
upper_crop_size = [rounding(x, 1) for x in abs(self.diff) / 2]
predicted = cropping(predicted, lower_crop_size, upper_crop_size)
else:
lower_pad_size = (abs(self.diff) // 2).tolist()
upper_pad_size = [rounding(x, 1) for x in abs(self.diff) / 2]
predicted = padding(predicted, lower_pad_size, upper_pad_size)
return predicted
def restore(self, predict_array_list):
predict_array = np.zeros_like(sitk.GetArrayFromImage(self.label))
size = np.array(self.label.GetSize()) - self.label_patch_size
indices = [
i for i in product(range(0, size[0] + 1, self.slide[0]),
range(0, size[1] + 1, self.slide[1]),
range(0, size[2] + 1, self.slide[2]))
]
with tqdm(total=len(predict_array_list),
desc="Restoring image...",
ncols=60) as pbar:
for pre_array, idx in zip(predict_array_list, indices):
x_slice = slice(idx[0], idx[0] + self.label_patch_size[0])
y_slice = slice(idx[1], idx[1] + self.label_patch_size[1])
z_slice = slice(idx[2], idx[2] + self.label_patch_size[2])
predict_array[z_slice, y_slice, x_slice] = pre_array
pbar.update(1)
predict = getImageWithMeta(predict_array, self.label)
predict = cropping(predict, self.lower_pad_size[1].tolist(),
self.upper_pad_size[1].tolist())
predict.SetOrigin(self.label.GetOrigin())
return predict
def main(args):
use_cuda = torch.cuda.is_available() and True
device = torch.device("cuda" if use_cuda else "cpu")
""" Slice module. """
labelFile = Path(args.imageDirectory) / "segmentation.nii.gz"
imageFile = Path(args.imageDirectory) / "imaging.nii.gz"
label = sitk.ReadImage(str(labelFile))
image = sitk.ReadImage(str(imageFile))
slicer = sler(image=image,
label=label,
outputImageSize=args.patchSize,
widthSize=args.widthSize,
paddingSize=args.paddingSize,
noFlip=args.noFlip)
slicer.execute()
_, cuttedImageArrayList = slicer.output("Array")
""" Load model. """
with open(args.modelweightfile, 'rb') as f:
model = cloudpickle.load(f)
model = torch.nn.DataParallel(model, device_ids=args.gpuid)
model.eval()
""" Segmentation module. """
segmentedArrayList = [[] for _ in range(2)]
for i in range(2):
length = len(cuttedImageArrayList[i])
for x in tqdm(range(length), desc="Segmenting images...", ncols=60):
imageArray = cuttedImageArrayList[i][x]
imageArray = imageArray.transpose((2, 0, 1))
imageArray = torch.from_numpy(imageArray[np.newaxis, ...]).to(
device, dtype=torch.float)
segmentedArray = model(imageArray)
segmentedArray = segmentedArray.to("cpu").detach().numpy().astype(
np.float)
segmentedArray = np.squeeze(segmentedArray)
segmentedArray = np.argmax(segmentedArray, axis=0).astype(np.uint8)
segmentedArrayList[i].append(segmentedArray)
""" Restore module. """
segmentedArray = slicer.restore(segmentedArrayList)
segmented = getImageWithMeta(segmentedArray, label)
createParentPath(args.savePath)
print("Saving image to {}".format(args.savePath))
sitk.WriteImage(segmented, args.savePath, True)
def main(args):
image = sitk.ReadImage(args.image_path)
image_array = sitk.GetArrayFromImage(image)
image_array_win = setWindowSize(image_array=image_array,
min_value=args.min_value,
max_value=args.max_value)
image_win = getImageWithMeta(image_array_win, image)
save_path = Path(args.save_path)
save_path.parent.mkdir(parents=True, exist_ok=True)
print("Saving images to {} ...".format(str(save_path)))
sitk.WriteImage(image_win, str(save_path), True)
print("Done.")
def main(args):
label = sitk.ReadImage(args.label_path)
label_array = sitk.GetArrayFromImage(label)
if args.mask_number < 0:
mask_array = (label_array > 0).astype(np.int)
else:
mask_array = (label_array == args.mask_number).astype(np.int)
mask = getImageWithMeta(mask_array, label)
createParentPath(args.save_path)
print("Saving mask image to {} ...".format(args.save_path))
sitk.WriteImage(mask, args.save_path, True)
print("Done")
def restore(self, predict_array_list):
predict_array = np.zeros_like(sitk.GetArrayFromImage(self.meta["padded_label"]))
with tqdm(total=len(predict_array_list), desc="Restoring image...", ncols=60) as pbar:
for pre_array, idx in zip(predict_array_list, self.meta["total_patch_idx"]):
z_slice = slice(idx[0], idx[0] + self.meta["patch_size"][0])
y_slice = slice(idx[1], idx[1] + self.meta["patch_size"][1])
x_slice = slice(idx[2], idx[2] + self.meta["patch_size"][2])
predict_array[x_slice, y_slice, z_slice] = pre_array
pbar.update(1)
predict = getImageWithMeta(predict_array, self.label)
predict = cropping(predict, self.meta["lower_padding_size"], self.meta["upper_padding_size"])
predict.SetOrigin(self.label.GetOrigin())
return predict
def restore(self, predicted_array_list):
predicted_array = np.zeros(self.padded_label.GetSize()[::-1])
z_size = self.padded_label.GetSize()[2] - self.label_patch_width
patch_size = np.array(self.padded_label.GetSize())
patch_size[2] = self.label_patch_width
indices = [i for i in range(0, z_size + 1, self.z_slide)]
with tqdm(total=len(indices), desc="Restoring image...",
ncols=60) as pbar:
for pre_array, index in zip(predicted_array_list, indices):
z_slice = slice(index, index + patch_size[2])
predicted_array[z_slice, ...] = pre_array
pbar.update(1)
predicted = getImageWithMeta(predicted_array, self.padded_label)
predicted = cropping(predicted, self.lower_pad_size[1].tolist(),
self.upper_pad_size[1].tolist())
return predicted
def execute(self):
self.meta = [{} for _ in range(2)]
labelArray = sitk.GetArrayFromImage(self.label)
imageArray = sitk.GetArrayFromImage(self.image)
matchobj = re.match("([0-9]+)-([0-9]+)-([0-9]+)", self.outputImageSize)
if matchobj is None:
print("[ERROR] Invalid patch size : {}".format(
self.outputImageSize))
sys.exit()
self.outputImageSize = [int(s) for s in matchobj.groups()]
self.outputLabelSize = self.outputImageSize[:2] + [1]
startIndex, endIndex = searchBound(labelArray, "sagittal")
if len(startIndex) != 2:
print("The patient has horse shoe kidney")
sys.exit()
largestKidneyLabelArray = []
largestKidneyImageArray = []
largestSlice = slice(endIndex[0], labelArray.shape[0])
largestKidneyLabelArray.append(labelArray[largestSlice, ...])
largestKidneyImageArray.append(imageArray[largestSlice, ...])
self.meta[0]["largestSlice"] = largestSlice
largestSlice = slice(0, startIndex[1])
largestKidneyLabelArray.append(labelArray[largestSlice, ...])
largestKidneyImageArray.append(imageArray[largestSlice, ...])
self.meta[1]["largestSlice"] = largestSlice
if not self.noFlip:
largestKidneyLabelArray[1] = largestKidneyLabelArray[1][::-1, ...]
largestKidneyImageArray[1] = largestKidneyImageArray[1][::-1, ...]
else:
print("The kidney doesn't flip.")
self.cuttedLabelArrayList = [[] for _ in range(2)]
self.cuttedStackedLabelArrayList = [[] for _ in range(2)]
self.cuttedImageArrayList = [[] for _ in range(2)]
for i in range(2):
p = (self.paddingSize, self.paddingSize)
largestKidneyLabelArray[i] = np.pad(largestKidneyLabelArray[i],
[p, p, (0, 0)], "minimum")
largestKidneyImageArray[i] = np.pad(largestKidneyImageArray[i],
[p, p, (0, 0)], "minimum")
startIndex, endIndex = caluculateClipSize(
largestKidneyLabelArray[i], "axial")
axialSlice = slice(startIndex, endIndex)
largestKidneyLabelArray[i] = largestKidneyLabelArray[i][...,
axialSlice]
largestKidneyImageArray[i] = largestKidneyImageArray[i][...,
axialSlice]
self.meta[i]["axialSlice"] = axialSlice
length = largestKidneyLabelArray[i].shape[2]
area = []
for x in range(length):
area.append(caluculateArea(largestKidneyLabelArray[i][..., x]))
maxArea = np.argmax(area)
maxAreaLabelArray = largestKidneyLabelArray[i][..., maxArea]
s, e = caluculateClipSize(maxAreaLabelArray[..., np.newaxis],
"sagittal")
width = e - s
s, e = caluculateClipSize(maxAreaLabelArray[..., np.newaxis],
"coronal")
height = e - s
wh = max(width, height)
margin = self.outputImageSize[2] // 2
top, bottom = caluculateClipSize(largestKidneyLabelArray[i],
"axial",
widthSize=0)
topSliceArray = largestKidneyLabelArray[i][..., top]
bottomSliceArray = largestKidneyLabelArray[i][..., bottom - 1]
check = False
sagittalSlices = []
coronalSlices = []
size = []
for x in tqdm(range(length), desc="Slicing images...", ncols=60):
a = caluculateArea(largestKidneyLabelArray[i][..., x])
if a == 0:
if not check:
sliceLabelArray = topSliceArray
else:
sliceLabelArray = bottomSliceArray
else:
sliceLabelArray = largestKidneyLabelArray[i][..., x]
check = True
center = getCenterOfGravity(sliceLabelArray)
x0 = center[0] - wh // 2
x1 = center[0] + wh // 2
y0 = center[1] - wh // 2
y1 = center[1] + wh // 2
minLabelArray = np.zeros(
(x1 - x0, y1 - y0)) + largestKidneyLabelArray[i].min()
minImageArray = np.zeros(
(x1 - x0, y1 - y0)) + largestKidneyImageArray[i].min()
sagittalSlice = slice(x0, x1)
coronalSlice = slice(y0, y1)
sagittalSlices.append(sagittalSlice)
coronalSlices.append(coronalSlice)
cuttedImageArray = []
cuttedStackedLabelArray = []
for y in range(-margin, margin + 1):
if 0 <= x + y < length:
cuttedStackedLabelArray.append(
largestKidneyLabelArray[i][sagittalSlice,
coronalSlice, x + y])
cuttedImageArray.append(
largestKidneyImageArray[i][sagittalSlice,
coronalSlice, x + y])
else:
cuttedStackedLabelArray.append(minLabelArray)
cuttedImageArray.append(minImageArray)
cuttedLabelArray = largestKidneyLabelArray[i][sagittalSlice,
coronalSlice, x]
cuttedStackedLabelArray = np.dstack(cuttedStackedLabelArray)
cuttedImageArray = np.dstack(cuttedImageArray)
self.cuttedLabelArrayList[i].append(cuttedLabelArray)
self.cuttedStackedLabelArrayList[i].append(
cuttedStackedLabelArray)
self.cuttedImageArrayList[i].append(cuttedImageArray)
size.append(cuttedLabelArray.shape)
self.meta[i]["sagittalSlice"] = sagittalSlices
self.meta[i]["coronalSlice"] = coronalSlices
self.meta[i]["size"] = size
""" For resampling, get direction, spacing, origin and minimun value in image and label in 2D. """
extractSliceFilter = sitk.ExtractImageFilter()
size = list(self.image.GetSize())
size[0] = 0
index = (0, 0, 0)
extractSliceFilter.SetSize(size)
extractSliceFilter.SetIndex(index)
self.sliceImage = extractSliceFilter.Execute(self.image)
self.cuttedLabelList = [[] for _ in range(2)]
self.cuttedImageList = [[] for _ in range(2)]
for i in range(2):
length = len(self.cuttedLabelArrayList[i])
for x in tqdm(range(length),
desc="Transforming images...",
ncols=60):
cuttedLabel = getImageWithMeta(
self.cuttedLabelArrayList[i][x][..., np.newaxis],
self.image)
cuttedImage = getImageWithMeta(self.cuttedImageArrayList[i][x],
self.image)
cuttedLabel = resampleSize(cuttedLabel,
self.outputLabelSize[::-1],
is_label=True)
cuttedImage = resampleSize(cuttedImage,
self.outputImageSize[::-1])
cuttedLabelArray = sitk.GetArrayFromImage(cuttedLabel)
cuttedImageArray = sitk.GetArrayFromImage(cuttedImage)
self.cuttedLabelArrayList[i][x] = cuttedLabelArray
self.cuttedImageArrayList[i][x] = cuttedImageArray
self.cuttedLabelList[i].append(cuttedLabel)
self.cuttedImageList[i].append(cuttedImage)
def main(_):
""" Slice module. """
labelFile = Path(args.imageDirectory) / "segmentation.nii.gz"
imageFile = Path(args.imageDirectory) / "imaging.nii.gz"
label = sitk.ReadImage(str(labelFile))
image = sitk.ReadImage(str(imageFile))
slicer = sler(image,
label,
outputImageSize=args.outputImageSize,
widthSize=args.widthSize,
paddingSize=args.paddingSize,
noFlip=args.noFlip)
slicer.execute()
_, cuttedImageArrayList = slicer.output("Array")
""" Segmentation module. """
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = True
sess = tf.compat.v1.Session(config=config)
tf.compat.v1.keras.backend.set_session(sess)
modelweightfile = os.path.expanduser(args.modelweightfile)
with tf.device('/device:GPU:{}'.format(args.gpuid)):
print('loading U-net model {}...'.format(modelweightfile),
end='',
flush=True)
#with open(args.modelfile) as f:
# model = tf.compat.v1.keras.models.model_from_yaml(f.read())
# model.load_weights(args.modelweightfile)
model = tf.compat.v1.keras.models.load_model(modelweightfile,
custom_objects={
'penalty_categorical':
penalty_categorical,
'kidney_dice':
kidney_dice,
'cancer_dice':
cancer_dice
})
print('done')
segmentedArrayList = [[] for _ in range(2)]
for i in range(2):
length = len(cuttedImageArrayList[i])
for x in tqdm(range(length), desc="Segmenting images...", ncols=60):
imageArray = cuttedImageArrayList[i][x]
imageArray = imageArray[np.newaxis, ...]
segmentedArray = model.predict(imageArray,
batch_size=args.batchsize,
verbose=0)
segmentedArray = np.squeeze(segmentedArray)
segmentedArray = np.argmax(segmentedArray,
axis=-1).astype(np.uint8)
print((segmentedArray > 0).any())
segmentedArrayList[i].append(segmentedArray)
""" Restore module. """
segmentedArray = slicer.restore(segmentedArrayList)
segmented = getImageWithMeta(segmentedArray, label)
createParentPath(args.savePath)
print("Saving image to {}".format(args.savePath))
sitk.WriteImage(segmented, args.savePath, True)
