|
|
import os, sys |
|
|
currentdir = os.path.dirname(os.path.realpath(__file__)) |
|
|
parentdir = os.path.dirname(currentdir) |
|
|
sys.path.append(parentdir) |
|
|
|
|
|
PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True' |
|
|
|
|
|
import tensorflow as tf |
|
|
import voxelmorph as vxm |
|
|
from voxelmorph.tf.modelio import LoadableModel, store_config_args |
|
|
|
|
|
|
|
|
class WeaklySupervised(LoadableModel): |
|
|
|
|
|
@store_config_args |
|
|
def __init__(self, inshape, all_labels: [list, tuple], nb_unet_features=None, int_steps=5, bidir=False, **kwargs): |
|
|
""" |
|
|
Parameters: |
|
|
inshape: Input shape. e.g. (192, 192, 192) |
|
|
all_labels: List of all labels included in training segmentations. |
|
|
hot_labels: List of labels to output as one-hot maps. |
|
|
nb_unet_features: Unet convolutional features. See VxmDense documentation for more information. |
|
|
int_steps: Number of flow integration steps. The warp is non-diffeomorphic when this value is 0. |
|
|
kwargs: Forwarded to the internal VxmDense model. |
|
|
""" |
|
|
|
|
|
fix_segm = tf.keras.Input((*inshape, len(all_labels)), name='fix_segmentations_input') |
|
|
mov_segm = tf.keras.Input((*inshape, len(all_labels)), name='mov_segmentations_input') |
|
|
|
|
|
mov_img = tf.keras.Input((*inshape, 1), name='mov_image_input') |
|
|
|
|
|
unet_input_model = tf.keras.Model(inputs=[mov_segm, fix_segm], outputs=[mov_segm, fix_segm]) |
|
|
|
|
|
vxm_model = vxm.networks.VxmDense(inshape=inshape, |
|
|
nb_unet_features=nb_unet_features, |
|
|
input_model=unet_input_model, |
|
|
int_steps=int_steps, |
|
|
bidir=bidir, **kwargs) |
|
|
|
|
|
pred_img = vxm.layers.SpatialTransformer(interp_method='linear', indexing='ij', name='pred_fix_img')( |
|
|
[mov_img, vxm_model.references.pos_flow]) |
|
|
|
|
|
inputs = [mov_segm, fix_segm, mov_img] |
|
|
outputs = [pred_img] + vxm_model.outputs |
|
|
|
|
|
self.references = LoadableModel.ReferenceContainer() |
|
|
self.references.pred_segm = vxm_model.outputs[0] |
|
|
self.references.pred_img = pred_img |
|
|
self.references.pos_flow = vxm_model.references.pos_flow |
|
|
|
|
|
super().__init__(inputs=inputs, outputs=outputs) |
|
|
|
|
|
def get_registration_model(self): |
|
|
return tf.keras.Model(self.inputs, self.references.pos_flow) |
|
|
|
|
|
def register(self, mov_img, mov_segm, fix_segm): |
|
|
return self.get_registration_model().predict([mov_segm, fix_segm, mov_img]) |
|
|
|
|
|
def apply_transform(self, mov_img, mov_segm, fix_segm, interp_method='linear'): |
|
|
warp_model = self.get_registration_model() |
|
|
img_input = tf.keras.Input(shape=mov_img.shape[1:], name='input_img') |
|
|
pred_img = vxm.layers.SpatialTransformer(interp_method=interp_method)([img_input, warp_model.output]) |
|
|
return tf.keras.Model(warp_model.inputs, pred_img).predict([mov_segm, fix_segm, mov_img]) |
|
|
|
|
|
|
