def __init__(self, moving_image_size, fixed_image_size, batch_size, scale=0.1):
self._batch_size = batch_size
self._scale = scale
self._moving_grid_ref = layer_util.get_reference_grid(
grid_size=moving_image_size
)
self._fixed_grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)
def test_get_reference_grid(self):
want = tf.constant(
np.array([[[[0, 0, 0], [0, 0, 1], [0, 0, 2]],
[[0, 1, 0], [0, 1, 1], [0, 1, 2]]]],
dtype=np.float32))
get = layer_util.get_reference_grid(grid_size=[1, 2, 3])
self.check_equal(want, get)
def __init__(
self,
image_size: tuple,
out_channels: int,
num_channel_initial: int,
extract_levels: List[int],
out_kernel_initializer: str,
out_activation: str,
name: str = "GlobalNet",
**kwargs,
):
"""
Image is encoded gradually, i from level 0 to E.
Then, a densely-connected layer outputs an affine
transformation.
:param image_size: tuple, such as (dim1, dim2, dim3)
:param out_channels: int, number of channels for the output
:param num_channel_initial: int, number of initial channels
:param extract_levels: list, which levels from net to extract
:param out_kernel_initializer: not used
:param out_activation: not used
:param name: name of the backbone.
:param kwargs: additional arguments.
"""
super().__init__(
image_size=image_size,
out_channels=out_channels,
num_channel_initial=num_channel_initial,
out_kernel_initializer=out_kernel_initializer,
out_activation=out_activation,
name=name,
**kwargs,
)
# save parameters
assert out_channels == 3
self._extract_levels = extract_levels
self._extract_max_level = max(self._extract_levels) # E
self.reference_grid = layer_util.get_reference_grid(image_size)
self.transform_initial = tf.constant_initializer(
value=list(np.eye(4, 3).reshape((-1))))
# init layer variables
num_channels = [
num_channel_initial * (2**level)
for level in range(self._extract_max_level + 1)
] # level 0 to E
self._downsample_blocks = [
layer.DownSampleResnetBlock(filters=num_channels[i],
kernel_size=7 if i == 0 else 3)
for i in range(self._extract_max_level)
] # level 0 to E-1
self._conv3d_block = layer.Conv3dBlock(
filters=num_channels[-1]) # level E
self._dense_layer = layer.Dense(
units=12, bias_initializer=self.transform_initial)
def __init__(self, fixed_image_size, **kwargs):
"""
:param fixed_image_size: shape = [f_dim1, f_dim2, f_dim3]
or [f_dim1, f_dim2, f_dim3, ch] with the last channel for features
:param kwargs:
"""
super(Warping, self).__init__(**kwargs)
self.grid_ref = tf.expand_dims(
layer_util.get_reference_grid(grid_size=fixed_image_size), axis=0
) # shape = (1, f_dim1, f_dim2, f_dim3, 3)
def test_get_reference_grid():
"""
Test get_reference_grid by confirming that it generates
a sample grid test case to is_equal_tf's tolerance level.
"""
want = tf.constant(
np.array(
[[[[0, 0, 0], [0, 0, 1], [0, 0, 2]],
[[0, 1, 0], [0, 1, 1], [0, 1, 2]]]],
dtype=np.float32,
))
get = layer_util.get_reference_grid(grid_size=[1, 2, 3])
assert is_equal_tf(want, get)
def __init__(
self,
moving_image_size: tuple,
fixed_image_size: tuple,
batch_size: int,
name: str = "RandomTransformation3D",
trainable: bool = False,
):
"""
Abstract class for image transformation.
:param moving_image_size: (m_dim1, m_dim2, m_dim3)
:param fixed_image_size: (f_dim1, f_dim2, f_dim3)
:param batch_size: size of mini-batch
:param name: name of layer
:param trainable: if this layer is trainable
"""
super().__init__(trainable=trainable, name=name)
self.moving_image_size = moving_image_size
self.fixed_image_size = fixed_image_size
self.batch_size = batch_size
self.moving_grid_ref = get_reference_grid(grid_size=moving_image_size)
self.fixed_grid_ref = get_reference_grid(grid_size=fixed_image_size)
def __init__(
self,
moving_image_size: Tuple[int, ...],
fixed_image_size: Tuple[int, ...],
batch_size: int,
name: str = "RandomTransformation3D",
trainable: bool = False,
):
"""
Abstract class for image transformation.
:param moving_image_size: (m_dim1, m_dim2, m_dim3)
:param fixed_image_size: (f_dim1, f_dim2, f_dim3)
:param batch_size: total number of samples consumed per step, over all devices.
:param name: name of layer
:param trainable: if this layer is trainable
"""
super().__init__(trainable=trainable, name=name)
self.moving_image_size = moving_image_size
self.fixed_image_size = fixed_image_size
self.batch_size = batch_size
self.moving_grid_ref = get_reference_grid(grid_size=moving_image_size)
self.fixed_grid_ref = get_reference_grid(grid_size=fixed_image_size)
def __init__(
self,
image_size,
out_channels,
num_channel_initial,
extract_levels,
out_kernel_initializer,
out_activation,
**kwargs,
):
"""
Image is encoded gradually, i from level 0 to E.
Then, a densely-connected layer outputs an affine
transformation.
:param out_channels: int, number of channels for the output
:param num_channel_initial: int, number of initial channels
:param extract_levels: list, which levels from net to extract
:param out_activation: str, activation at last layer
:param out_kernel_initializer: str, which kernel to use as initialiser
:param kwargs:
"""
super(GlobalNet, self).__init__(**kwargs)
# save parameters
self._extract_levels = extract_levels
self._extract_max_level = max(self._extract_levels) # E
self.reference_grid = layer_util.get_reference_grid(image_size)
self.transform_initial = tf.constant_initializer(
value=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]
)
# init layer variables
num_channels = [
num_channel_initial * (2 ** level)
for level in range(self._extract_max_level + 1)
] # level 0 to E
self._downsample_blocks = [
layer.DownSampleResnetBlock(
filters=num_channels[i], kernel_size=7 if i == 0 else 3
)
for i in range(self._extract_max_level)
] # level 0 to E-1
self._conv3d_block = layer.Conv3dBlock(filters=num_channels[-1]) # level E
self._dense_layer = layer.Dense(
units=12, bias_initializer=self.transform_initial
)
self._reshape = tf.keras.layers.Reshape(target_shape=(4, 3))
def __init__(self,
fixed_image_size: tuple,
name: str = "warping",
**kwargs):
"""
Init.
:param fixed_image_size: shape = (f_dim1, f_dim2, f_dim3)
or (f_dim1, f_dim2, f_dim3, ch) with the last channel for features
:param name: name of the layer
:param kwargs: additional arguments.
"""
super().__init__(name=name, **kwargs)
self._fixed_image_size = fixed_image_size
# shape = (1, f_dim1, f_dim2, f_dim3, 3)
self.grid_ref = layer_util.get_reference_grid(
grid_size=fixed_image_size)[None, ...]
def __init__(
self,
image_size: tuple,
name: str = "AffineHead",
):
"""
Init.
:param image_size: such as (dim1, dim2, dim3)
:param name: name of the layer
"""
super().__init__(name=name)
self.reference_grid = layer_util.get_reference_grid(image_size)
self.transform_initial = tf.constant_initializer(
value=list(np.eye(4, 3).reshape((-1))))
self._flatten = tfkl.Flatten()
self._dense = tfkl.Dense(units=12,
bias_initializer=self.transform_initial)
def __init__(self, fixed_image_size: tuple, **kwargs):
"""
A layer warps an image using DDF.
Reference:
- transform of neuron
https://github.com/adalca/neurite/blob/legacy/neuron/utils.py
where vol = image, loc_shift = ddf
:param fixed_image_size: shape = (f_dim1, f_dim2, f_dim3)
or (f_dim1, f_dim2, f_dim3, ch) with the last channel for features
:param kwargs: additional arguments.
"""
super().__init__(**kwargs)
self.grid_ref = tf.expand_dims(
layer_util.get_reference_grid(grid_size=fixed_image_size), axis=0
) # shape = (1, f_dim1, f_dim2, f_dim3, 3)
def __init__(
self,
moving_image_size: Tuple,
fixed_image_size: Tuple,
index_size: int,
labeled: bool,
batch_size: int,
config: dict,
name: str = "RegistrationModel",
):
"""
Init.
:param moving_image_size: (m_dim1, m_dim2, m_dim3)
:param fixed_image_size: (f_dim1, f_dim2, f_dim3)
:param index_size: number of indices for identify each sample
:param labeled: if the data is labeled
:param batch_size: total number of samples consumed per step, over all devices.
When using multiple devices, TensorFlow automatically split the tensors.
Therefore, input shapes should be defined over batch_size.
:param config: config for method, backbone, and loss.
:param name: name of the model
"""
super().__init__(name=name)
self.moving_image_size = moving_image_size
self.fixed_image_size = fixed_image_size
self.index_size = index_size
self.labeled = labeled
self.config = config
self.batch_size = batch_size
self._inputs = None # save inputs of self._model as dict
self._outputs = None # save outputs of self._model as dict
self.grid_ref = layer_util.get_reference_grid(
grid_size=fixed_image_size)[None, ...]
self._model: tf.keras.Model = self.build_model()
self.build_loss()
def predict(
gpu: str,
gpu_allow_growth: bool,
ckpt_path: str,
mode: str,
batch_size: int,
log_dir: str,
sample_label: str,
config_path: (str, list),
save_nifti: bool = True,
save_png: bool = True,
):
"""
Function to predict some metrics from the saved model and logging results.
:param gpu: str, which env gpu to use.
:param gpu_allow_growth: bool, whether to allow gpu growth or not
:param ckpt_path: str, where model is stored, should be like log_folder/save/xxx.ckpt
:param mode: train / valid / test, to define which split of dataset to be evaluated
:param batch_size: int, batch size to perform predictions in
:param log_dir: str, path to store logs
:param sample_label: sample/all, not used
:param save_nifti: if true, outputs will be saved in nifti format
:param save_png: if true, outputs will be saved in png format
:param config_path: to overwrite the default config
"""
# TODO support custom sample_label
logging.warning(
"sample_label is not used in predict. It is True if and only if mode == 'train'."
)
# env vars
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
os.environ[
"TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true"
# load config
config, log_dir = build_config(config_path=config_path,
log_dir=log_dir,
ckpt_path=ckpt_path)
preprocess_config = config["train"]["preprocess"]
preprocess_config["batch_size"] = batch_size
# data
data_loader, dataset, _ = build_dataset(
dataset_config=config["dataset"],
preprocess_config=preprocess_config,
mode=mode,
training=False,
repeat=False,
)
# optimizer
optimizer = opt.build_optimizer(
optimizer_config=config["train"]["optimizer"])
# model
model = build_model(
moving_image_size=data_loader.moving_image_shape,
fixed_image_size=data_loader.fixed_image_shape,
index_size=data_loader.num_indices,
labeled=config["dataset"]["labeled"],
batch_size=preprocess_config["batch_size"],
model_config=config["train"]["model"],
loss_config=config["train"]["loss"],
)
# metrics
model.compile(optimizer=optimizer)
# load weights
# https://stackoverflow.com/questions/58289342/tf2-0-translation-model-error-when-restoring-the-saved-model-unresolved-objec
model.load_weights(ckpt_path).expect_partial()
# predict
fixed_grid_ref = tf.expand_dims(
layer_util.get_reference_grid(grid_size=data_loader.fixed_image_shape),
axis=0) # shape = (1, f_dim1, f_dim2, f_dim3, 3)
predict_on_dataset(
dataset=dataset,
fixed_grid_ref=fixed_grid_ref,
model=model,
model_method=config["train"]["model"]["method"],
save_dir=log_dir + "/test",
save_nifti=save_nifti,
save_png=save_png,
)
# close the opened files in data loaders
data_loader.close()
## load image
if not os.path.exists(DATA_PATH):
raise ("Download the data using demo_data.py script")
if not os.path.exists(FILE_PATH):
raise ("Download the data using demo_data.py script")
fid = h5py.File(FILE_PATH, "r")
fixed_image = tf.cast(tf.expand_dims(fid["image"], axis=0), dtype=tf.float32)
fixed_image = (fixed_image - tf.reduce_min(fixed_image)) / (
tf.reduce_max(fixed_image) - tf.reduce_min(fixed_image)
) # normalisation to [0,1]
# generate a radomly-affine-transformed moving image
fixed_image_size = fixed_image.shape
transform_random = layer_util.random_transform_generator(batch_size=1, scale=0.2)
grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size[1:4])
grid_random = layer_util.warp_grid(grid_ref, transform_random)
moving_image = layer_util.resample(vol=fixed_image, loc=grid_random)
# warp the labels to get ground-truth using the same random affine, for validation
fixed_labels = tf.cast(tf.expand_dims(fid["label"], axis=0), dtype=tf.float32)
moving_labels = tf.stack(
[
layer_util.resample(vol=fixed_labels[..., idx], loc=grid_random)
for idx in range(fixed_labels.shape[4])
],
axis=4,
)
## optimisation
@tf.function
def main(gpu, gpu_allow_growth, ckpt_path, mode, batch_size, log):
# sanity check
if not ckpt_path.endswith(
".ckpt"): # should be like log_folder/save/xxx.ckpt
raise ValueError("checkpoint path should end with .ckpt")
# env vars
os.environ['CUDA_VISIBLE_DEVICES'] = gpu
os.environ[
"TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true"
# load config
config = config_parser.load("/".join(ckpt_path.split("/")[:-2]) +
"/config.yaml")
data_config = config["data"]
if data_config["name"] == "mr_us":
data_config["sample_label"]["train"] = "all"
data_config["sample_label"]["test"] = "all"
tf_data_config = config["tf"]["data"]
tf_data_config["batch_size"] = batch_size
tf_opt_config = config["tf"]["opt"]
tf_model_config = config["tf"]["model"]
tf_loss_config = config["tf"]["loss"]
log_folder_name = log if log != "" else datetime.now().strftime(
"%Y%m%d-%H%M%S")
log_dir = config["log_dir"][:-1] if config["log_dir"][
-1] == "/" else config["log_dir"]
log_dir = log_dir + "/" + log_folder_name
# data
data_loader = load.get_data_loader(data_config, mode)
dataset = data_loader.get_dataset_and_preprocess(training=False,
repeat=False,
**tf_data_config)
# optimizer
optimizer = opt.get_optimizer(tf_opt_config)
# model
model = network.build_model(
moving_image_size=data_loader.moving_image_shape,
fixed_image_size=data_loader.fixed_image_shape,
index_size=data_loader.num_indices,
batch_size=tf_data_config["batch_size"],
tf_model_config=tf_model_config,
tf_loss_config=tf_loss_config)
# metrics
model.compile(optimizer=optimizer,
loss=label_loss.get_similarity_fn(
config=tf_loss_config["similarity"]["label"]),
metrics=[
metric.MeanDiceScore(),
metric.MeanCentroidDistance(
grid_size=data_loader.fixed_image_shape)
])
# load weights
model.load_weights(ckpt_path)
# predict
fixed_grid_ref = layer_util.get_reference_grid(
grid_size=data_loader.fixed_image_shape)
predict(data_loader=data_loader,
dataset=dataset,
fixed_grid_ref=fixed_grid_ref,
model=model,
save_dir=log_dir + "/test")
def predict(
gpu,
gpu_allow_growth,
ckpt_path,
mode,
batch_size,
log_dir,
sample_label,
config_path,
):
"""
Function to predict some metrics from the saved model and logging results.
:param gpu: str, which env gpu to use.
:param gpu_allow_growth: bool, whether to allow gpu growth or not
:param ckpt_path: str, where model is stored, should be like
log_folder/save/xxx.ckpt
:param mode: which mode to load the data ??
:param batch_size: int, batch size to perform predictions in
:param log_dir: str, path to store logs
:param sample_label:
:param config_path: to overwrite the default config
"""
logging.error("TODO sample_label is not used in predict")
# env vars
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
os.environ[
"TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true"
# load config
config, log_dir = init(log_dir, ckpt_path, config_path)
dataset_config = config["dataset"]
preprocess_config = config["train"]["preprocess"]
preprocess_config["batch_size"] = batch_size
optimizer_config = config["train"]["optimizer"]
model_config = config["train"]["model"]
loss_config = config["train"]["loss"]
# data
data_loader = load.get_data_loader(dataset_config, mode)
if data_loader is None:
raise ValueError(
"Data loader for prediction is None. Probably the data dir path is not defined."
)
dataset = data_loader.get_dataset_and_preprocess(training=False,
repeat=False,
**preprocess_config)
# optimizer
optimizer = opt.build_optimizer(optimizer_config)
# model
model = build_model(
moving_image_size=data_loader.moving_image_shape,
fixed_image_size=data_loader.fixed_image_shape,
index_size=data_loader.num_indices,
labeled=dataset_config["labeled"],
batch_size=preprocess_config["batch_size"],
model_config=model_config,
loss_config=loss_config,
)
# metrics
model.compile(optimizer=optimizer)
# load weights
# https://stackoverflow.com/questions/58289342/tf2-0-translation-model-error-when-restoring-the-saved-model-unresolved-objec
model.load_weights(ckpt_path).expect_partial()
# predict
fixed_grid_ref = layer_util.get_reference_grid(
grid_size=data_loader.fixed_image_shape)
predict_on_dataset(
dataset=dataset,
fixed_grid_ref=fixed_grid_ref,
model=model,
save_dir=log_dir + "/test",
)
data_loader.close()
def predict(
gpu: str,
gpu_allow_growth: bool,
ckpt_path: str,
mode: str,
batch_size: int,
exp_name: str,
config_path: Union[str, List[str]],
save_nifti: bool = True,
save_png: bool = True,
log_dir: str = "logs",
):
"""
Function to predict some metrics from the saved model and logging results.
:param gpu: which env gpu to use.
:param gpu_allow_growth: whether to allow gpu growth or not
:param ckpt_path: where model is stored, should be like log_folder/save/ckpt-x
:param mode: train / valid / test, to define which split of dataset to be evaluated
:param batch_size: int, batch size to perform predictions in
:param exp_name: name of the experiment
:param log_dir: path of the log directory
:param save_nifti: if true, outputs will be saved in nifti format
:param save_png: if true, outputs will be saved in png format
:param config_path: to overwrite the default config
"""
# TODO support custom sample_label
logging.warning("sample_label is not used in predict. "
"It is True if and only if mode == 'train'.")
# env vars
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
os.environ[
"TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true"
# load config
config, log_dir, ckpt_path = build_config(config_path=config_path,
log_dir=log_dir,
exp_name=exp_name,
ckpt_path=ckpt_path)
preprocess_config = config["train"]["preprocess"]
# batch_size corresponds to batch_size per GPU
gpus = tf.config.experimental.list_physical_devices("GPU")
preprocess_config["batch_size"] = batch_size * max(len(gpus), 1)
# data
data_loader, dataset, _ = build_dataset(
dataset_config=config["dataset"],
preprocess_config=preprocess_config,
mode=mode,
training=False,
repeat=False,
)
assert data_loader is not None
# optimizer
optimizer = opt.build_optimizer(
optimizer_config=config["train"]["optimizer"])
# model
model: tf.keras.Model = REGISTRY.build_model(config=dict(
name=config["train"]["method"],
moving_image_size=data_loader.moving_image_shape,
fixed_image_size=data_loader.fixed_image_shape,
index_size=data_loader.num_indices,
labeled=config["dataset"]["labeled"],
batch_size=config["train"]["preprocess"]["batch_size"],
config=config["train"],
))
# metrics
model.compile(optimizer=optimizer)
# load weights
if ckpt_path.endswith(".ckpt"):
# for ckpt from tf.keras.callbacks.ModelCheckpoint
# skip warnings because of optimizers
# https://stackoverflow.com/questions/58289342/tf2-0-translation-model-error-when-restoring-the-saved-model-unresolved-object
model.load_weights(ckpt_path).expect_partial() # pragma: no cover
else:
# for ckpts from ckpt manager callback
_, _ = build_checkpoint_callback(
model=model,
dataset=dataset,
log_dir=log_dir,
save_period=config["train"]["save_period"],
ckpt_path=ckpt_path,
)
# predict
fixed_grid_ref = tf.expand_dims(
layer_util.get_reference_grid(grid_size=data_loader.fixed_image_shape),
axis=0) # shape = (1, f_dim1, f_dim2, f_dim3, 3)
predict_on_dataset(
dataset=dataset,
fixed_grid_ref=fixed_grid_ref,
model=model,
model_method=config["train"]["method"],
save_dir=os.path.join(log_dir, "test"),
save_nifti=save_nifti,
save_png=save_png,
)
# close the opened files in data loaders
data_loader.close()
def predict(
gpu: str,
ckpt_path: str,
split: str,
batch_size: int,
exp_name: str,
config_path: Union[str, List[str]],
num_workers: int = 1,
gpu_allow_growth: bool = True,
save_nifti: bool = True,
save_png: bool = True,
log_dir: str = "logs",
):
"""
Function to predict some metrics from the saved model and logging results.
:param gpu: which env gpu to use.
:param ckpt_path: where model is stored, should be like log_folder/save/ckpt-x.
:param split: train / valid / test, to define the split to be evaluated.
:param batch_size: int, batch size to perform predictions.
:param exp_name: name of the experiment.
:param config_path: to overwrite the default config.
:param num_workers: number of cpu cores to be used, <=0 means not limited.
:param gpu_allow_growth: whether to allocate whole GPU memory for training.
:param save_nifti: if true, outputs will be saved in nifti format.
:param save_png: if true, outputs will be saved in png format.
:param log_dir: path of the log directory.
"""
# env vars
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
os.environ[
"TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true"
if num_workers <= 0: # pragma: no cover
logger.info(
"Limiting CPU usage by setting environment variables "
"OMP_NUM_THREADS, TF_NUM_INTRAOP_THREADS, TF_NUM_INTEROP_THREADS to %d. "
"This may slow down the prediction. "
"Please use --num_workers flag to modify the behavior. "
"Setting to 0 or negative values will remove the limitation.",
num_workers,
)
# limit CPU usage
# https://github.com/tensorflow/tensorflow/issues/29968#issuecomment-789604232
os.environ["OMP_NUM_THREADS"] = str(num_workers)
os.environ["TF_NUM_INTRAOP_THREADS"] = str(num_workers)
os.environ["TF_NUM_INTEROP_THREADS"] = str(num_workers)
# load config
config, log_dir, ckpt_path = build_config(config_path=config_path,
log_dir=log_dir,
exp_name=exp_name,
ckpt_path=ckpt_path)
config["train"]["preprocess"]["batch_size"] = batch_size
# data
data_loader, dataset, _ = build_dataset(
dataset_config=config["dataset"],
preprocess_config=config["train"]["preprocess"],
split=split,
training=False,
repeat=False,
)
assert data_loader is not None
# use strategy to support multiple GPUs
# the network is mirrored in each GPU so that we can use larger batch size
# https://www.tensorflow.org/guide/distributed_training
# only model, optimizer and metrics need to be defined inside the strategy
num_devices = max(len(tf.config.list_physical_devices("GPU")), 1)
if num_devices > 1: # pragma: no cover
strategy = tf.distribute.MirroredStrategy()
if batch_size % num_devices != 0:
raise ValueError(
f"batch size {batch_size} can not be divided evenly "
f"by the number of devices.")
else:
strategy = tf.distribute.get_strategy()
with strategy.scope():
model: tf.keras.Model = REGISTRY.build_model(config=dict(
name=config["train"]["method"],
moving_image_size=data_loader.moving_image_shape,
fixed_image_size=data_loader.fixed_image_shape,
index_size=data_loader.num_indices,
labeled=config["dataset"][split]["labeled"],
batch_size=batch_size,
config=config["train"],
))
optimizer = opt.build_optimizer(
optimizer_config=config["train"]["optimizer"])
model.compile(optimizer=optimizer)
model.plot_model(output_dir=log_dir)
# load weights
if ckpt_path.endswith(".ckpt"):
# for ckpt from tf.keras.callbacks.ModelCheckpoint
# skip warnings because of optimizers
# https://stackoverflow.com/questions/58289342/tf2-0-translation-model-error-when-restoring-the-saved-model-unresolved-object
model.load_weights(ckpt_path).expect_partial() # pragma: no cover
else:
# for ckpts from ckpt manager callback
_, _ = build_checkpoint_callback(
model=model,
dataset=dataset,
log_dir=log_dir,
save_period=config["train"]["save_period"],
ckpt_path=ckpt_path,
)
# predict
fixed_grid_ref = tf.expand_dims(
layer_util.get_reference_grid(grid_size=data_loader.fixed_image_shape),
axis=0) # shape = (1, f_dim1, f_dim2, f_dim3, 3)
predict_on_dataset(
dataset=dataset,
fixed_grid_ref=fixed_grid_ref,
model=model,
save_dir=os.path.join(log_dir, "test"),
save_nifti=save_nifti,
save_png=save_png,
)
# close the opened files in data loaders
data_loader.close()
def __init__(self, grid_size, name="metric/centroid_distance_mean", **kwargs):
super(MeanCentroidDistance, self).__init__(name=name, **kwargs)
self.grid = layer_util.get_reference_grid(grid_size)
