def save(self, save_path, patient_id, with_nonmask=False):
if not isinstance(patient_id, str):
patient_id = str(patient_id)
save_path = Path(save_path)
save_mask_path = save_path / "mask" / "case_{}".format(
patient_id.zfill(2))
save_mask_path.mkdir(parents=True, exist_ok=True)
if with_nonmask:
save_nonmask_path = save_path / "nonmask" / "case_{}".format(
patient_id.zfill(2))
save_nonmask_path.mkdir(parents=True, exist_ok=True)
if with_nonmask:
desc = "Saving masked and nonmasked images, labels and coordinates..."
else:
desc = "Saving masked images, labels and coordinates..."
with tqdm(total=self.image_patch_array_generator.__len__(),
ncols=100,
desc=desc) as pbar:
for i, (ipa, tpa, lpa, cpa, mpa, _,
_) in enumerate(self.generateData()):
if isMasked(mpa):
save_masked_image_path = save_mask_path / "image_{:04d}.npy".format(
i)
save_masked_trans_path = save_mask_path / "trans_{:04d}.npy".format(
i)
save_masked_label_path = save_mask_path / "label_{:04d}.npy".format(
i)
save_masked_coord_path = save_mask_path / "coord_{:04d}.npy".format(
i)
np.save(str(save_masked_image_path), ipa)
np.save(str(save_masked_trans_path), tpa)
np.save(str(save_masked_label_path), lpa)
np.save(str(save_masked_coord_path), cpa)
else:
if with_nonmask:
save_nonmasked_image_path = save_nonmask_path / "image_{:04d}.npy".format(
i)
save_nonmasked_trans_path = save_nonmask_path / "trans_{:04d}.npy".format(
i)
save_nonmasked_label_path = save_nonmask_path / "label_{:04d}.npy".format(
i)
save_nonmasked_coord_path = save_nonmask_path / "coord_{:04d}.npy".format(
i)
np.save(str(save_nonmasked_image_path), ipa)
np.save(str(save_nonmasked_trans_path), tpa)
np.save(str(save_nonmasked_label_path), lpa)
np.save(str(save_nonmasked_coord_path), cpa)
pbar.update(1)
def savePatchArray(self,
save_dir,
patient_id,
input_name="input",
target_name="target",
with_nonmask=False):
""" Save patch array to save_path.
save_path : save_dir / 'all' or 'mask' or 'nonmask' / case_{patient_id} / 'input_name_0000.npy' or 'target_name_0000.npy'
Parameters:
save_dir (str) -- Refer above.
patient_id (str) -- Refer above.
input_name (str) -- Input patch array name. Refer above.
target_name (str) -- Target patch array name. Refer above.
with_nonmask (bool) -- Whether you save nonmasked images when mask image is fed.
"""
save_dir = Path(save_dir)
if self.mask is None:
save_mask_path = save_dir / "all" / "case_{}".format(
str(patient_id).zfill(2)
) # When mask_image is None, all of patch arrays are masked.
else:
save_mask_path = save_dir / "mask" / "case_{}".format(
str(patient_id).zfill(2))
save_nonmask_path = save_dir / "nonmask" / "case_{}".format(
str(patient_id).zfill(2))
save_mask_path.mkdir(parents=True, exist_ok=True)
if with_nonmask:
save_nonmask_path.mkdir(parents=True, exist_ok=True)
with tqdm(total=self.image_generator.__len__(),
desc="Savig patch array...",
ncols=60) as pbar:
for i, (ipa, tpa, mpa, _) in enumerate(self.generatePatchArray()):
if isMasked(mpa):
input_masked_save_path = save_mask_path / "input_{:04d}.npy".format(
i)
target_masked_save_path = save_mask_path / "target_{:04d}.npy".format(
i)
np.save(str(input_masked_save_path), ipa)
np.save(str(target_masked_save_path), tpa)
else:
if with_nonmask:
input_nonmasked_save_path = save_nonmasked_path / "input_{:04d}.npy".format(
i)
target_nonmasked_save_path = save_nonmasked_path / "target_{:04d}.npy".format(
i)
np.save(str(input_nonmasked_save_path), ipa)
np.save(str(target_nonmasked_sace_path), tpa)
pbar.update(1)
def main(args):
printArgs(args)
""" Read images. """
image = sitk.ReadImage(args.image_path)
mask = sitkReadImageElseNone(args.mask_path)
""" Dummy image for prediction"""
label = sitk.Image(image.GetSize(), sitk.sitkInt8)
label.SetOrigin(image.GetOrigin())
label.SetDirection(image.GetDirection())
label.SetSpacing(image.GetSpacing())
""" Get the patch size from string."""
image_patch_size = getSizeFromStringElseNone(args.image_patch_size)
label_patch_size = getSizeFromStringElseNone(args.label_patch_size)
patch_generator = Extractor(image=image,
label=label,
mask=mask,
image_patch_size=image_patch_size,
label_patch_size=label_patch_size,
overlap=args.overlap,
num_class=args.num_class,
class_axis=args.axis)
""" Load model. """
with open(args.modelweightfile, 'rb') as f:
model = pickle5.load(f)
# Use a single gpu because It does't work when we use multi gpu.
model = torch.nn.DataParallel(model, device_ids=[args.gpuid[0]])
model.eval()
""" Segmentation module. """
use_cuda = torch.cuda.is_available() and True
device = torch.device("cuda" if use_cuda else "cpu")
segmenter = Segmenter(model, device=device)
transformer = transform()
with tqdm(total=patch_generator.__len__(),
ncols=60,
desc="Segmenting and restoring...") as pbar:
for image_patch, _, mask_patch, _, index in patch_generator.generateData(
):
image_patch_array, mask_patch_array = transformer(
"test", image_patch, mask_patch)
if isMasked(mask_patch_array):
segmented_array = segmenter(image_patch_array)
#segmented_array = np.argmax(segmented_array, axis=0)
patch_generator.insertToPredictedArray(index, segmented_array)
pbar.update(1)
segmented = patch_generator.outputRestoredImage()
save_path = Path(args.save_path)
save_path.parent.mkdir(parents=True, exist_ok=True)
print("Saving image to {}".format(str(save_path)))
sitk.WriteImage(segmented, str(save_path), True)
def save(self, save_path, patient_id, with_nonmask=False):
if not isinstance(patient_id, str):
patient_id = str(patient_id)
save_path = Path(save_path)
save_mask_path = save_path / "mask" / "case_{}".format(
patient_id.zfill(2))
save_mask_path.mkdir(parents=True, exist_ok=True)
if with_nonmask:
save_nonmask_path = save_path / "nonmask" / "case_{}".format(
patient_id.zfill(2))
save_nonmask_path.mkdir(parents=True, exist_ok=True)
desc = "Saving masked and nonmasked images and labels..."
else:
desc = "Saving masked images and labels..."
with tqdm(total=self.image_patch_generator.__len__(),
ncols=100,
desc=desc) as pbar:
for i, (ipa, lpa, mpa, _, _) in enumerate(self.generateData()):
if isMasked(mpa):
save_masked_image_path = save_mask_path / "image_{:04d}.mha".format(
i)
save_masked_label_path = save_mask_path / "label_{:04d}.mha".format(
i)
sitk.WriteImage(ipa, str(save_masked_image_path), True)
sitk.WriteImage(lpa, str(save_masked_label_path), True)
else:
if with_nonmask:
save_nonmasked_image_path = save_nonmask_path / "image_{:04d}.mha".format(
i)
save_nonmasked_label_path = save_nonmask_path / "label_{:04d}.mha".format(
i)
sitk.WriteImage(ipa, str(save_nonmasked_image_path),
True)
sitk.WriteImage(lpa, str(save_nonmasked_label_path),
True)
pbar.update(1)
def main(args):
printArgs(args)
""" Read images. """
image = sitk.ReadImage(args.image_path)
trans = sitk.ReadImage(args.trans_path)
liver = sitk.ReadImage(args.liver_path)
if args.mask_path is not None:
mask = sitk.ReadImage(args.mask_path)
else:
mask = None
""" Dummy image for prediction"""
label = sitk.Image(image.GetSize(), sitk.sitkInt8)
label.SetOrigin(image.GetOrigin())
label.SetDirection(image.GetDirection())
label.SetSpacing(image.GetSpacing())
""" Get the patch size from string."""
image_patch_size = getSizeFromString(args.image_patch_size)
label_patch_size = getSizeFromString(args.label_patch_size)
cogc = CenterOfGravityCalculater(liver)
liver_center = cogc.execute()
print("Liver center: ", liver_center)
iace = ImageAndCoordinateExtractor(
image=image,
trans=trans,
label=label,
mask=mask,
image_array_patch_size=image_patch_size,
label_array_patch_size=label_patch_size,
overlap=args.overlap,
center=liver_center,
)
""" 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. """
use_cuda = torch.cuda.is_available() and True
device = torch.device("cuda" if use_cuda else "cpu")
""" For pre-processing and post-processing for segmentation. """
predictor = Predictor(model, device=device)
with tqdm(total=iace.__len__(),
ncols=60,
desc="Segmenting and restoring...") as pbar:
for image_patch_array, trans_patch_array, label_patch_array, coord_patch_array, mask_patch_array, _, index in iace.generateData(
):
for rate in args.rate_list:
transform = UNetTransform(test_mix_rate=rate)
if isMasked(mask_patch_array):
input_array_list = [
image_patch_array, trans_patch_array, coord_patch_array
]
input_array_list, _ = transform("test", input_array_list,
label_patch_array)
segmented_array = predictor(input_array_list)
iace.insertToPredictedArray(index, segmented_array)
pbar.update(1)
segmented = iace.outputRestoredImage()
save_path = Path(args.save_path)
save_path.parent.mkdir(parents=True, exist_ok=True)
print("Saving image to {}".format(str(save_path)))
sitk.WriteImage(segmented, str(save_path), True)