def test_3d_4d(self, arr: Tuple[tf.Tensor, np.ndarray]): save_array(save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True) assert self.get_num_files_in_dir(self.png_dir, suffix=".png") == 4 assert self.get_num_files_in_dir(self.save_dir, suffix=".nii.gz") == 1
def test_wrong_shape(self, arr: (tf.Tensor, np.ndarray), err_msg: str): with pytest.raises(ValueError) as err_info: save_array(save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True) assert err_msg in str(err_info.value)
def test_wrong_shape(self, arr: (tf.Tensor, np.ndarray), err_msg: str): """test TensorFlow/Numpy inputs with incorrect shapes""" with pytest.raises(ValueError) as err_info: save_array(save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True) assert err_msg in str(err_info.value)
def test_3d_4d(self, arr: (tf.Tensor, np.ndarray)): """test 3d/4d TensorFlow/Numpy inputs""" save_array(save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True) assert self.get_num_files_in_dir(self.png_dir, suffix=".png") == 4 assert self.get_num_files_in_dir(self.save_dir, suffix=".nii.gz") == 1
def test_save_png(self, save_png: bool): arr = np.random.rand(2, 3, 4, 3) save_array( save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True, save_png=save_png, ) assert (self.get_num_files_in_dir(self.png_dir, suffix=".png") == int(save_png) * 4)
def test_save_nifti(self, save_nifti: bool): arr = np.random.rand(2, 3, 4, 3) save_array( save_dir=self.save_dir, arr=arr, name=self.arr_name, normalize=True, save_nifti=save_nifti, ) assert self.get_num_files_in_dir(self.save_dir, suffix=".nii.gz") == int(save_nifti)
def test_overwrite(self, overwrite: bool): arr1 = np.random.rand(2, 3, 4, 3) arr2 = arr1 + 1 nifti_file_path = os.path.join(self.save_dir, self.arr_name + ".nii.gz") # save arr1 os.makedirs(self.save_dir, exist_ok=True) nib.save(img=nib.Nifti1Image(arr1, affine=np.eye(4)), filename=nifti_file_path) # save arr2 w/o overwrite save_array( save_dir=self.save_dir, arr=arr2, name=self.arr_name, normalize=True, overwrite=overwrite, ) arr_read = load_nifti_file(file_path=nifti_file_path) assert is_equal_np(arr2 if overwrite else arr1, arr_read)
def predict_on_dataset( dataset: tf.data.Dataset, fixed_grid_ref: tf.Tensor, model: tf.keras.Model, model_method: str, save_dir: str, save_nifti: bool, save_png: bool, ): """ Function to predict results from a dataset from some model :param dataset: where data is stored :param fixed_grid_ref: shape=(1, f_dim1, f_dim2, f_dim3, 3) :param model: model to be used for prediction :param model_method: str, ddf / dvf / affine / conditional :param save_dir: str, path to store dir :param save_nifti: if true, outputs will be saved in nifti format :param save_png: if true, outputs will be saved in png format """ # remove the save_dir in case it exists if os.path.exists(save_dir): shutil.rmtree(save_dir) sample_index_strs = [] metric_lists = [] for _, inputs_dict in enumerate(dataset): outputs_dict = model.predict(x=inputs_dict) # moving image/label # (batch, m_dim1, m_dim2, m_dim3) moving_image = inputs_dict["moving_image"] moving_label = inputs_dict.get("moving_label", None) # fixed image/labelimage_index # (batch, f_dim1, f_dim2, f_dim3) fixed_image = inputs_dict["fixed_image"] fixed_label = inputs_dict.get("fixed_label", None) # indices to identify the pair # (batch, num_indices) last indice is for label, -1 means unlabeled data indices = inputs_dict.get("indices") # ddf / dvf # (batch, f_dim1, f_dim2, f_dim3, 3) ddf = outputs_dict.get("ddf", None) dvf = outputs_dict.get("dvf", None) affine = outputs_dict.get("affine", None) # (batch, 4, 3) # prediction # (batch, f_dim1, f_dim2, f_dim3) pred_fixed_label = outputs_dict.get("pred_fixed_label", None) pred_fixed_image = (layer_util.resample( vol=moving_image, loc=fixed_grid_ref + ddf) if ddf is not None else None) # save images of inputs and outputs for sample_index in range(moving_image.shape[0]): # save moving/fixed image under pair_dir # save moving/fixed label, pred fixed image/label, ddf/dvf under label dir # if labeled, label dir is a sub dir of pair_dir, otherwise = pair_dir # init output path indices_i = indices[sample_index, :].numpy().astype(int).tolist() pair_dir, label_dir = build_pair_output_path(indices=indices_i, save_dir=save_dir) # save image/label # if model is conditional, the pred_fixed_image depends on the input label conditional = model_method == "conditional" arr_save_dirs = [ pair_dir, pair_dir, label_dir if conditional else pair_dir, label_dir, label_dir, label_dir, ] arrs = [ moving_image, fixed_image, pred_fixed_image, moving_label, fixed_label, pred_fixed_label, ] names = [ "moving_image", "fixed_image", "pred_fixed_image", # or warped moving image "moving_label", "fixed_label", "pred_fixed_label", # or warped moving label ] for arr_save_dir, arr, name in zip(arr_save_dirs, arrs, names): if arr is not None: # for files under pair_dir, do not overwrite save_array( save_dir=arr_save_dir, arr=arr[sample_index, :, :, :], name=name, gray=True, save_nifti=save_nifti, save_png=save_png, overwrite=arr_save_dir == label_dir, ) # save ddf / dvf arrs = [ddf, dvf] names = ["ddf", "dvf"] for arr, name in zip(arrs, names): if arr is not None: arr = normalize_array(arr=arr[sample_index, :, :, :]) save_array( save_dir=label_dir if conditional else pair_dir, arr=arr, name=name, gray=False, save_nifti=save_nifti, save_png=save_png, ) # save affine if affine is not None: np.savetxt( fname=os.path.join(label_dir if conditional else pair_dir, "affine.txt"), x=affine[sample_index, :, :].numpy(), delimiter=",", ) # calculate metric sample_index_str = "_".join([str(x) for x in indices_i]) if sample_index_str in sample_index_strs: raise ValueError( "Sample is repeated, maybe the dataset has been repeated.") sample_index_strs.append(sample_index_str) metric = calculate_metrics( fixed_image=fixed_image, fixed_label=fixed_label, pred_fixed_image=pred_fixed_image, pred_fixed_label=pred_fixed_label, fixed_grid_ref=fixed_grid_ref, sample_index=sample_index, ) metric["pair_index"] = indices_i[:-1] metric["label_index"] = indices_i[-1] metric_lists.append(metric) # save metric save_metric_dict(save_dir=save_dir, metrics=metric_lists)
shutil.rmtree(SAVE_PATH) os.mkdir(SAVE_PATH) arrays = [ tf.squeeze(a) for a in [ moving_image, fixed_image, warped_moving_image, moving_label, fixed_label, warped_moving_label, var_ddf, ] ] arr_names = [ "moving_image", "fixed_image", "warped_moving_image", "moving_label", "fixed_label", "warped_moving_label", "ddf", ] for arr, arr_name in zip(arrays, arr_names): util.save_array( save_dir=SAVE_PATH, arr=arr, name=arr_name, normalize=True, save_png=False ) os.chdir(MAIN_PATH)
arrays = [ tf.transpose(a, [1, 2, 3, 0]) if a.ndim == 4 else tf.squeeze(a) for a in [ moving_image, fixed_image, warped_moving_image, moving_labels, fixed_labels, warped_moving_labels, ] ] arr_names = [ "moving_image", "fixed_image", "warped_moving_image", "moving_label", "fixed_label", "warped_moving_label", ] for arr, arr_name in zip(arrays, arr_names): for n in range(arr.shape[-1]): util.save_array( save_dir=SAVE_PATH, arr=arr[..., n], name=arr_name + (arr.shape[-1] > 1) * "_{}".format(n), normalize="image" in arr_name, # label's value is already in [0, 1] ) os.chdir(MAIN_PATH)
os.mkdir(SAVE_PATH) arrays = [ tf.transpose(a, [1, 2, 3, 0]) if a.ndim == 4 else tf.squeeze(a) for a in [ moving_image, fixed_image, warped_moving_image, moving_labels, fixed_labels, warped_moving_labels, ] ] arr_names = [ "moving_image", "fixed_image", "warped_moving_image", "moving_label", "fixed_label", "warped_moving_label", ] for arr, arr_name in zip(arrays, arr_names): for n in range(arr.shape[-1]): util.save_array( save_dir=SAVE_PATH, arr=arr[..., n], name=arr_name + (arr.shape[-1] > 1) * "_{}".format(n), gray=True, ) os.chdir(MAIN_PATH)
def test_save_array(): """ Test save_array by testing different shapes and count output files """ def get_num_pngs_in_dir(dir_paths): return len([x for x in os.listdir(dir_paths) if x.endswith(".png")]) def get_num_niftis_in_dir(dir_paths): return len([x for x in os.listdir(dir_paths) if x.endswith(".nii.gz")]) save_dir = "logs/test_util_save_array" if os.path.exists(save_dir): shutil.rmtree(save_dir) # test 3D tf tensor name = "3d_tf" out_dir = os.path.join(save_dir, name) arr = tf.random.uniform(shape=(2, 3, 4)) save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert get_num_pngs_in_dir(out_dir) == 4 assert get_num_niftis_in_dir(save_dir) == 1 shutil.rmtree(out_dir) os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 4D tf tensor name = "4d_tf" out_dir = os.path.join(save_dir, name) arr = tf.random.uniform(shape=(2, 3, 4, 3)) save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert get_num_pngs_in_dir(out_dir) == 4 assert get_num_niftis_in_dir(save_dir) == 1 shutil.rmtree(out_dir) os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 3D np tensor name = "3d_np" out_dir = os.path.join(save_dir, name) arr = np.random.rand(2, 3, 4) save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert get_num_pngs_in_dir(out_dir) == 4 assert get_num_niftis_in_dir(save_dir) == 1 shutil.rmtree(out_dir) os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 4D np tensor name = "4d_np" out_dir = os.path.join(save_dir, name) arr = np.random.rand(2, 3, 4, 3) save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert get_num_pngs_in_dir(out_dir) == 4 assert get_num_niftis_in_dir(save_dir) == 1 shutil.rmtree(out_dir) os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 4D np tensor without nifti name = "4d_np" out_dir = os.path.join(save_dir, name) arr = np.random.rand(2, 3, 4, 3) save_array(save_dir=save_dir, arr=arr, name=name, gray=True, save_nifti=False) assert get_num_pngs_in_dir(out_dir) == 4 assert get_num_niftis_in_dir(save_dir) == 0 shutil.rmtree(out_dir) # test 4D np tensor without png name = "4d_np" out_dir = os.path.join(save_dir, name) arr = np.random.rand(2, 3, 4, 3) assert not os.path.exists(out_dir) save_array(save_dir=save_dir, arr=arr, name=name, gray=True, save_png=False) assert not os.path.exists(out_dir) assert get_num_niftis_in_dir(save_dir) == 1 os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 4D np tensor with overwrite name = "4d_np" out_dir = os.path.join(save_dir, name) arr1 = np.random.rand(2, 3, 4, 3) arr2 = np.random.rand(2, 3, 4, 3) assert not is_equal_np(arr1, arr2) nifti_file_path = os.path.join(save_dir, name + ".nii.gz") # save arr1 os.makedirs(save_dir, exist_ok=True) nib.save(img=nib.Nifti2Image(arr1, affine=np.eye(4)), filename=nifti_file_path) # save arr2 without overwrite save_array(save_dir=save_dir, arr=arr1, name=name, gray=True, overwrite=False) arr_read = load_nifti_file(file_path=nifti_file_path) assert is_equal_np(arr1, arr_read) # save arr2 with overwrite save_array(save_dir=save_dir, arr=arr2, name=name, gray=True, overwrite=True) arr_read = load_nifti_file(file_path=nifti_file_path) assert is_equal_np(arr2, arr_read) shutil.rmtree(out_dir) os.remove(os.path.join(save_dir, name + ".nii.gz")) # test 5D np tensor name = "5d_np" arr = np.random.rand(2, 3, 4, 1, 3) with pytest.raises(ValueError) as err_info: save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert "arr must be 3d or 4d numpy array or tf tensor" in str( err_info.value) # test 4D np tensor with wrong shape name = "5d_np" arr = np.random.rand(2, 3, 4, 1) with pytest.raises(ValueError) as err_info: save_array(save_dir=save_dir, arr=arr, name=name, gray=True) assert "4d arr must have 3 channels as last dimension" in str( err_info.value)
def predict_on_dataset( dataset: tf.data.Dataset, fixed_grid_ref: tf.Tensor, model: tf.keras.Model, model_method: str, save_dir: str, save_nifti: bool, save_png: bool, ): """ Function to predict results from a dataset from some model :param dataset: where data is stored :param fixed_grid_ref: shape=(1, f_dim1, f_dim2, f_dim3, 3) :param model: model to be used for prediction :param model_method: ddf / dvf / affine / conditional :param save_dir: path to store dir :param save_nifti: if true, outputs will be saved in nifti format :param save_png: if true, outputs will be saved in png format """ # remove the save_dir in case it exists if os.path.exists(save_dir): shutil.rmtree(save_dir) # pragma: no cover sample_index_strs = [] metric_lists = [] for _, inputs in enumerate(dataset): batch_size = inputs[list(inputs.keys())[0]].shape[0] outputs = model.predict(x=inputs, batch_size=batch_size) indices, processed = model.postprocess(inputs=inputs, outputs=outputs) # convert to np arrays indices = indices.numpy() processed = { k: (v[0].numpy() if isinstance(v[0], tf.Tensor) else v[0], v[1], v[2]) for k, v in processed.items() } # save images of inputs and outputs for sample_index in range(batch_size): # save label independent tensors under pair_dir, otherwise under label_dir # init output path indices_i = indices[sample_index, :].astype(int).tolist() pair_dir, label_dir = build_pair_output_path(indices=indices_i, save_dir=save_dir) for name, (arr, normalize, on_label) in processed.items(): if name == "theta": np.savetxt( fname=os.path.join(pair_dir, "affine.txt"), X=arr[sample_index, :, :], delimiter=",", ) continue arr_save_dir = label_dir if on_label else pair_dir save_array( save_dir=arr_save_dir, arr=arr[sample_index, :, :, :], name=name, normalize=normalize, # label's value is already in [0, 1] save_nifti=save_nifti, save_png=save_png, overwrite=arr_save_dir == label_dir, ) # calculate metric sample_index_str = "_".join([str(x) for x in indices_i]) if sample_index_str in sample_index_strs: # pragma: no cover raise ValueError( "Sample is repeated, maybe the dataset has been repeated.") sample_index_strs.append(sample_index_str) metric = calculate_metrics( fixed_image=processed["fixed_image"][0], fixed_label=processed["fixed_label"][0] if model.labeled else None, pred_fixed_image=processed["pred_fixed_image"][0], pred_fixed_label=processed["pred_fixed_label"][0] if model.labeled else None, fixed_grid_ref=fixed_grid_ref, sample_index=sample_index, ) metric["pair_index"] = indices_i[:-1] metric["label_index"] = indices_i[-1] metric_lists.append(metric) # save metric save_metric_dict(save_dir=save_dir, metrics=metric_lists)