
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
import keras
from keras import layers, models
from keras import initializers
from keras.models import Model
from keras.layers import Conv2D, MaxPooling2D, Input, Conv2DTranspose, Concatenate, BatchNormalization, UpSampling2D
from keras.utils.conv_utils import conv_output_length, deconv_length
from keras.layers import Dropout, Activation
from keras.optimizers import Adam, SGD
from keras.layers.advanced_activations import LeakyReLU
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
from keras import backend as K
from keras.utils import plot_model
import tensorflow as tf
import glob
import random
#import cv2
from random import shuffle


def dice_soft(y_true, y_pred, loss_type='sorensen', axis=[1, 2, 3], smooth=1e-5, from_logits=False):
    if not from_logits:
        _epsilon = tf.convert_to_tensor(1e-7, y_pred.dtype.base_dtype)
        y_pred = tf.clip_by_value(y_pred, _epsilon, 1 - _epsilon)
        y_pred = tf.log(y_pred / (1 - y_pred))

    inse = tf.reduce_sum(y_pred * y_true, axis=axis)
    if loss_type == 'jaccard':
        l = tf.reduce_sum(y_pred * y_pred, axis=axis)
        r = tf.reduce_sum(y_true * y_true, axis=axis)
    elif loss_type == 'sorensen':
        l = tf.reduce_sum(y_pred, axis=axis)
        r = tf.reduce_sum(y_true, axis=axis)
    else:
        raise Exception("Unknow loss_type")
    dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    dice = tf.reduce_mean(dice)
    return dice


def dice_hard(y_true, y_pred, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    y_pred = tf.cast(y_pred > threshold, dtype=tf.float32)
    y_true = tf.cast(y_true > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(y_pred, y_true), axis=axis)
    l = tf.reduce_sum(y_pred, axis=axis)
    r = tf.reduce_sum(y_true, axis=axis)
    hard_dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    hard_dice = tf.reduce_mean(hard_dice)
    return hard_dice


def dice_loss(y_true, y_pred, from_logits=False):
    return 1 - dice_soft(y_true, y_pred)

def CapsNetR3(input_shape, n_class=2):
    x = layers.Input(shape=input_shape)

    # Layer 1: Just a conventional Conv2D layer
    conv1 = layers.Conv2D(filters=16, kernel_size=5, strides=1, padding='same', activation='relu', name='conv1')(x)

    # Reshape layer to be 1 capsule x [filters] atoms
    _, H, W, C = conv1.get_shape()
    conv1_reshaped = layers.Reshape((H.value, W.value, 1, C.value))(conv1)

    # Layer 1: Primary Capsule: Conv cap with routing 1
    primary_caps = ConvCapsuleLayer(kernel_size=5, num_capsule=2, num_atoms=16, strides=2, padding='same',
                                    routings=1, name='primarycaps')(conv1_reshaped)

    # Layer 2: Convolutional Capsule
    conv_cap_2_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=16, strides=1, padding='same',
                                    routings=3, name='conv_cap_2_1')(primary_caps)

    # Layer 2: Convolutional Capsule
    conv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=32, strides=2, padding='same',
                                    routings=3, name='conv_cap_2_2')(conv_cap_2_1)

    # Layer 3: Convolutional Capsule
    conv_cap_3_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=32, strides=1, padding='same',
                                    routings=3, name='conv_cap_3_1')(conv_cap_2_2)

    # Layer 3: Convolutional Capsule
    conv_cap_3_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=64, strides=2, padding='same',
                                    routings=3, name='conv_cap_3_2')(conv_cap_3_1)

    # Layer 4: Convolutional Capsule
    conv_cap_4_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=32, strides=1, padding='same',
                                    routings=3, name='conv_cap_4_1')(conv_cap_3_2)

    # Layer 1 Up: Deconvolutional Capsule
    deconv_cap_1_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=8, num_atoms=32, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_1_1')(conv_cap_4_1)

    # Skip connection
    up_1 = layers.Concatenate(axis=-2, name='up_1')([deconv_cap_1_1, conv_cap_3_1])

    # Layer 1 Up: Deconvolutional Capsule
    deconv_cap_1_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=32, strides=1,
                                      padding='same', routings=3, name='deconv_cap_1_2')(up_1)

    # Layer 2 Up: Deconvolutional Capsule
    deconv_cap_2_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=4, num_atoms=16, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_2_1')(deconv_cap_1_2)

    # Skip connection
    up_2 = layers.Concatenate(axis=-2, name='up_2')([deconv_cap_2_1, conv_cap_2_1])

    # Layer 2 Up: Deconvolutional Capsule
    deconv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=16, strides=1,
                                      padding='same', routings=3, name='deconv_cap_2_2')(up_2)

    # Layer 3 Up: Deconvolutional Capsule
    deconv_cap_3_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=2, num_atoms=16, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_3_1')(deconv_cap_2_2)

    # Skip connection
    up_3 = layers.Concatenate(axis=-2, name='up_3')([deconv_cap_3_1, conv1_reshaped])

    # Layer 4: Convolutional Capsule: 1x1
    seg_caps = ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=16, strides=1, padding='same',
                                routings=3, name='seg_caps')(up_3)

    # Layer 4: This is an auxiliary layer to replace each capsule with its length. Just to match the true label's shape.
    out_seg = Length(num_classes=n_class, seg=True, name='out_seg')(seg_caps)

    # Decoder network.
    _, H, W, C, A = seg_caps.get_shape()
    y = layers.Input(shape=input_shape[:-1] + (1,))
    masked_by_y = Mask()([seg_caps, y])  # The true label is used to mask the output of capsule layer. For training
    masked = Mask()(seg_caps)  # Mask using the capsule with maximal length. For prediction

    def shared_decoder(mask_layer):
        recon_remove_dim = layers.Reshape((H.value, W.value, A.value))(mask_layer)

        recon_1 = layers.Conv2D(filters=64, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                activation='relu', name='recon_1')(recon_remove_dim)

        recon_2 = layers.Conv2D(filters=128, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                activation='relu', name='recon_2')(recon_1)

        out_recon = layers.Conv2D(filters=1, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                  activation='sigmoid', name='out_recon')(recon_2)

        return out_recon

    # Models for training and evaluation (prediction)
    train_model = models.Model(inputs=[x, y], outputs=[out_seg, shared_decoder(masked_by_y)])
    eval_model = models.Model(inputs=x, outputs=[out_seg, shared_decoder(masked)])

    # manipulate model
    noise = layers.Input(shape=((H.value, W.value, C.value, A.value)))
    noised_seg_caps = layers.Add()([seg_caps, noise])
    masked_noised_y = Mask()([noised_seg_caps, y])
    manipulate_model = models.Model(inputs=[x, y, noise], outputs=shared_decoder(masked_noised_y))

    # return train_model, eval_model, manipulate_model

    return train_model


class Length(layers.Layer):
    def __init__(self, num_classes, seg=True, **kwargs):
        super(Length, self).__init__(**kwargs)
        if num_classes == 2:
            self.num_classes = 1
        else:
            self.num_classes = num_classes
        self.seg = seg

    def call(self, inputs, **kwargs):
        if inputs.get_shape().ndims == 5:
            assert inputs.get_shape()[-2].value == 1, 'Error: Must have num_capsules = 1 going into Length'
            inputs = K.squeeze(inputs, axis=-2)
        return K.expand_dims(tf.norm(inputs, axis=-1), axis=-1)

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 5:
            input_shape = input_shape[0:-2] + input_shape[-1:]
        if self.seg:
            return input_shape[:-1] + (self.num_classes,)
        else:
            return input_shape[:-1]

    def get_config(self):
        config = {'num_classes': self.num_classes, 'seg': self.seg}
        base_config = super(Length, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class Mask(layers.Layer):
    def __init__(self, resize_masks=False, **kwargs):
        super(Mask, self).__init__(**kwargs)
        self.resize_masks = resize_masks

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input

        else:
            if inputs.get_shape().ndims == 3:
                x = K.sqrt(K.sum(K.square(inputs), -1))
                mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.get_shape().as_list()[1])
                masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
            else:
                masked = inputs

        return masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:  # true label provided
            if len(input_shape[0]) == 3:
                return tuple([None, input_shape[0][1] * input_shape[0][2]])
            else:
                return input_shape[0]
        else:  # no true label provided
            if len(input_shape) == 3:
                return tuple([None, input_shape[1] * input_shape[2]])
            else:
                return input_shape

    def get_config(self):
        config = {'resize_masks': self.resize_masks}
        base_config = super(Mask, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConvCapsuleLayer(layers.Layer):
    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1, padding='same', routings=3,
                 kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 5, "The input Tensor should have shape=[None, input_height, input_width," \
                                      " input_num_capsule, input_num_atoms]"
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]

        # Transform matrix
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                        self.input_num_atoms, self.num_capsule * self.num_atoms],
                                 initializer=self.kernel_initializer,
                                 name='W')

        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.num_atoms],
                                 initializer=initializers.constant(0.1),
                                 name='b')

        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [
            input_shape[0] * input_shape[1], self.input_height, self.input_width, self.input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.input_width, self.input_num_atoms))

        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.strides),
                        padding=self.padding, data_format='channels_last')

        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()

        votes = K.reshape(conv, [input_shape[1], input_shape[0], votes_shape[1], votes_shape[2],
                                 self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value, conv_width.value,
                         self.num_capsule, self.num_atoms))

        logit_shape = K.stack([
            input_shape[1], input_shape[0], votes_shape[1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.value, 1, 1])

        activations = update_routing(
            votes=votes,
            biases=biases_replicated,
            logit_shape=logit_shape,
            num_dims=6,
            input_dim=self.input_num_capsule,
            output_dim=self.num_capsule,
            num_routing=self.routings)

        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(
                space[i],
                self.kernel_size,
                padding=self.padding,
                stride=self.strides,
                dilation=1)
            new_space.append(new_dim)

        return (input_shape[0],) + tuple(new_space) + (self.num_capsule, self.num_atoms)

    def get_config(self):
        config = {
            'kernel_size': self.kernel_size,
            'num_capsule': self.num_capsule,
            'num_atoms': self.num_atoms,
            'strides': self.strides,
            'padding': self.padding,
            'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.kernel_initializer)
        }
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):
    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2, upsamp_type='deconv', padding='same', routings=3,
                 kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 5, "The input Tensor should have shape=[None, input_height, input_width," \
                                      " input_num_capsule, input_num_atoms]"
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]

        # Transform matrix
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.input_num_atoms,
                                            self.num_capsule * self.num_atoms * self.scaling * self.scaling],
                                     initializer=self.kernel_initializer,
                                     name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.input_num_atoms, self.num_capsule * self.num_atoms],
                                     initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.num_capsule * self.num_atoms, self.input_num_atoms],
                                     initializer=self.kernel_initializer, name='W')
        else:
            raise NotImplementedError('Upsampling must be one of: "deconv", "resize", or "subpix"')

        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.num_atoms],
                                 initializer=initializers.constant(0.1),
                                 name='b')

        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [
            input_shape[1] * input_shape[0], self.input_height, self.input_width, self.input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.input_width, self.input_num_atoms))

        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling, self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1), padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(1, 1), padding='same',
                            data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]

            # Infer the dynamic output shape:
            out_height = deconv_length(self.input_height, self.scaling, self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.num_capsule * self.num_atoms)

            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W, output_shape, (self.scaling, self.scaling),
                                         padding=self.padding, data_format='channels_last')

        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()

        votes = K.reshape(outputs, [input_shape[1], input_shape[0], votes_shape[1], votes_shape[2],
                                    self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value, conv_width.value,
                         self.num_capsule, self.num_atoms))

        logit_shape = K.stack([
            input_shape[1], input_shape[0], votes_shape[1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2], 1, 1])

        activations = update_routing(
            votes=votes,
            biases=biases_replicated,
            logit_shape=logit_shape,
            num_dims=6,
            input_dim=self.input_num_capsule,
            output_dim=self.num_capsule,
            num_routing=self.routings)

        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)

        output_shape[1] = deconv_length(output_shape[1], self.scaling, self.kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self.kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms

        return tuple(output_shape)

    def get_config(self):
        config = {
            'kernel_size': self.kernel_size,
            'num_capsule': self.num_capsule,
            'num_atoms': self.num_atoms,
            'scaling': self.scaling,
            'padding': self.padding,
            'upsamp_type': self.upsamp_type,
            'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.kernel_initializer)
        }
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def update_routing(votes, biases, logit_shape, num_dims, input_dim, output_dim,
                   num_routing):
    if num_dims == 6:
        votes_t_shape = [5, 0, 1, 2, 3, 4]
        r_t_shape = [1, 2, 3, 4, 5, 0]
    elif num_dims == 4:
        votes_t_shape = [3, 0, 1, 2]
        r_t_shape = [1, 2, 3, 0]
    else:
        raise NotImplementedError('Not implemented')

    votes_trans = tf.transpose(votes, votes_t_shape)
    _, _, _, height, width, caps = votes_trans.get_shape()

    def _body(i, logits, activations):
        """Routing while loop."""
        # route: [batch, input_dim, output_dim, ...]
        route = tf.nn.softmax(logits, dim=-1)
        preactivate_unrolled = route * votes_trans
        preact_trans = tf.transpose(preactivate_unrolled, r_t_shape)
        preactivate = tf.reduce_sum(preact_trans, axis=1) + biases
        activation = _squash(preactivate)
        activations = activations.write(i, activation)
        act_3d = K.expand_dims(activation, 1)
        tile_shape = np.ones(num_dims, dtype=np.int32).tolist()
        tile_shape[1] = input_dim
        act_replicated = tf.tile(act_3d, tile_shape)
        distances = tf.reduce_sum(votes * act_replicated, axis=-1)
        logits += distances
        return (i + 1, logits, activations)

    activations = tf.TensorArray(
        dtype=tf.float32, size=num_routing, clear_after_read=False)
    logits = tf.fill(logit_shape, 0.0)

    i = tf.constant(0, dtype=tf.int32)
    _, logits, activations = tf.while_loop(
        lambda i, logits, activations: i < num_routing,
        _body,
        loop_vars=[i, logits, activations],
        swap_memory=True)

    return K.cast(activations.read(num_routing - 1), dtype='float32')

def _squash(input_tensor):
    norm = tf.norm(input_tensor, axis=-1, keep_dims=True)
    norm_squared = norm * norm
    return (input_tensor / norm) * (norm_squared / (1 + norm_squared))
def get_model():
    model = CapsNetR3((256, 256, 1))
    opt = Adam(lr=3e-5)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics={'out_seg': dice_hard})
    model.load_weights('weights/croppedtransferlesionsegcaps.h5')
    return model

def main_predict(model, path,truemaskfile):
    raw = Image.open(path)
    raw = np.array(raw.resize((256, 256))) / 255.
    raw=raw[:,:,0]
    raw = np.expand_dims(raw, 2)
    print(raw.shape)
    true_msk = Image.open(truemaskfile.replace('vol', 'lesion'))
    true_msk = np.array(true_msk.resize((256, 256))) / 255.
    true_msk = np.expand_dims(true_msk, 2)
    print(true_msk.shape)
    pred = model.predict([np.expand_dims(raw, 0), np.expand_dims(true_msk, 0)])
    msk = pred[0].squeeze()
    msk = np.stack((msk,) * 3, axis=-1)
    print(np.unique(msk))
    msk[msk >= 0.5] = 1
    msk[msk < 0.5] = 0
    return msk


4ba017d62b84601ce4d3e425f5454e7ce90840c8
<mask token>


class Mask(layers.Layer):
    <mask token>

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input
        elif inputs.get_shape().ndims == 3:
            x = K.sqrt(K.sum(K.square(inputs), -1))
            mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.
                get_shape().as_list()[1])
            masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
        else:
            masked = inputs
        return masked
    <mask token>
    <mask token>


class ConvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1,
        padding='same', routings=3, kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
            self.input_num_atoms, self.num_capsule * self.num_atoms],
            initializer=self.kernel_initializer, name='W')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[0] *
            input_shape[1], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.
            strides), padding=self.padding, data_format='channels_last')
        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()
        votes = K.reshape(conv, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.
            value, 1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(space[i], self.kernel_size,
                padding=self.padding, stride=self.strides, dilation=1)
            new_space.append(new_dim)
        return (input_shape[0],) + tuple(new_space) + (self.num_capsule,
            self.num_atoms)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'strides': self.
            strides, 'padding': self.padding, 'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.
            kernel_initializer)}
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2,
        upsamp_type='deconv', padding='same', routings=3,
        kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms * self.scaling * self.scaling], initializer=self.
                kernel_initializer, name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms], initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.num_capsule * self.num_atoms, self.
                input_num_atoms], initializer=self.kernel_initializer, name='W'
                )
        else:
            raise NotImplementedError(
                'Upsampling must be one of: "deconv", "resize", or "subpix"')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[1] *
            input_shape[0], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling,
                self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1),
                padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(
                1, 1), padding='same', data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]
            out_height = deconv_length(self.input_height, self.scaling,
                self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.
                kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.
                num_capsule * self.num_atoms)
            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W,
                output_shape, (self.scaling, self.scaling), padding=self.
                padding, data_format='channels_last')
        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()
        votes = K.reshape(outputs, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2],
            1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        output_shape[1] = deconv_length(output_shape[1], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms
        return tuple(output_shape)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'scaling': self.
            scaling, 'padding': self.padding, 'upsamp_type': self.
            upsamp_type, 'routings': self.routings, 'kernel_initializer':
            initializers.serialize(self.kernel_initializer)}
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


<mask token>

<mask token>


class Length(layers.Layer):

    def __init__(self, num_classes, seg=True, **kwargs):
        super(Length, self).__init__(**kwargs)
        if num_classes == 2:
            self.num_classes = 1
        else:
            self.num_classes = num_classes
        self.seg = seg

    def call(self, inputs, **kwargs):
        if inputs.get_shape().ndims == 5:
            assert inputs.get_shape()[-2
                ].value == 1, 'Error: Must have num_capsules = 1 going into Length'
            inputs = K.squeeze(inputs, axis=-2)
        return K.expand_dims(tf.norm(inputs, axis=-1), axis=-1)

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 5:
            input_shape = input_shape[0:-2] + input_shape[-1:]
        if self.seg:
            return input_shape[:-1] + (self.num_classes,)
        else:
            return input_shape[:-1]

    def get_config(self):
        config = {'num_classes': self.num_classes, 'seg': self.seg}
        base_config = super(Length, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class Mask(layers.Layer):

    def __init__(self, resize_masks=False, **kwargs):
        super(Mask, self).__init__(**kwargs)
        self.resize_masks = resize_masks

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input
        elif inputs.get_shape().ndims == 3:
            x = K.sqrt(K.sum(K.square(inputs), -1))
            mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.
                get_shape().as_list()[1])
            masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
        else:
            masked = inputs
        return masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:
            if len(input_shape[0]) == 3:
                return tuple([None, input_shape[0][1] * input_shape[0][2]])
            else:
                return input_shape[0]
        elif len(input_shape) == 3:
            return tuple([None, input_shape[1] * input_shape[2]])
        else:
            return input_shape

    def get_config(self):
        config = {'resize_masks': self.resize_masks}
        base_config = super(Mask, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1,
        padding='same', routings=3, kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
            self.input_num_atoms, self.num_capsule * self.num_atoms],
            initializer=self.kernel_initializer, name='W')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[0] *
            input_shape[1], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.
            strides), padding=self.padding, data_format='channels_last')
        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()
        votes = K.reshape(conv, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.
            value, 1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(space[i], self.kernel_size,
                padding=self.padding, stride=self.strides, dilation=1)
            new_space.append(new_dim)
        return (input_shape[0],) + tuple(new_space) + (self.num_capsule,
            self.num_atoms)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'strides': self.
            strides, 'padding': self.padding, 'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.
            kernel_initializer)}
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2,
        upsamp_type='deconv', padding='same', routings=3,
        kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms * self.scaling * self.scaling], initializer=self.
                kernel_initializer, name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms], initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.num_capsule * self.num_atoms, self.
                input_num_atoms], initializer=self.kernel_initializer, name='W'
                )
        else:
            raise NotImplementedError(
                'Upsampling must be one of: "deconv", "resize", or "subpix"')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[1] *
            input_shape[0], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling,
                self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1),
                padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(
                1, 1), padding='same', data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]
            out_height = deconv_length(self.input_height, self.scaling,
                self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.
                kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.
                num_capsule * self.num_atoms)
            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W,
                output_shape, (self.scaling, self.scaling), padding=self.
                padding, data_format='channels_last')
        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()
        votes = K.reshape(outputs, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2],
            1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        output_shape[1] = deconv_length(output_shape[1], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms
        return tuple(output_shape)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'scaling': self.
            scaling, 'padding': self.padding, 'upsamp_type': self.
            upsamp_type, 'routings': self.routings, 'kernel_initializer':
            initializers.serialize(self.kernel_initializer)}
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


<mask token>

<mask token>


def dice_loss(y_true, y_pred, from_logits=False):
    return 1 - dice_soft(y_true, y_pred)


def CapsNetR3(input_shape, n_class=2):
    x = layers.Input(shape=input_shape)
    conv1 = layers.Conv2D(filters=16, kernel_size=5, strides=1, padding=
        'same', activation='relu', name='conv1')(x)
    _, H, W, C = conv1.get_shape()
    conv1_reshaped = layers.Reshape((H.value, W.value, 1, C.value))(conv1)
    primary_caps = ConvCapsuleLayer(kernel_size=5, num_capsule=2, num_atoms
        =16, strides=2, padding='same', routings=1, name='primarycaps')(
        conv1_reshaped)
    conv_cap_2_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms
        =16, strides=1, padding='same', routings=3, name='conv_cap_2_1')(
        primary_caps)
    conv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms
        =32, strides=2, padding='same', routings=3, name='conv_cap_2_2')(
        conv_cap_2_1)
    conv_cap_3_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =32, strides=1, padding='same', routings=3, name='conv_cap_3_1')(
        conv_cap_2_2)
    conv_cap_3_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =64, strides=2, padding='same', routings=3, name='conv_cap_3_2')(
        conv_cap_3_1)
    conv_cap_4_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =32, strides=1, padding='same', routings=3, name='conv_cap_4_1')(
        conv_cap_3_2)
    deconv_cap_1_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=8,
        num_atoms=32, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_1_1')(conv_cap_4_1)
    up_1 = layers.Concatenate(axis=-2, name='up_1')([deconv_cap_1_1,
        conv_cap_3_1])
    deconv_cap_1_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4,
        num_atoms=32, strides=1, padding='same', routings=3, name=
        'deconv_cap_1_2')(up_1)
    deconv_cap_2_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=4,
        num_atoms=16, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_2_1')(deconv_cap_1_2)
    up_2 = layers.Concatenate(axis=-2, name='up_2')([deconv_cap_2_1,
        conv_cap_2_1])
    deconv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4,
        num_atoms=16, strides=1, padding='same', routings=3, name=
        'deconv_cap_2_2')(up_2)
    deconv_cap_3_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=2,
        num_atoms=16, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_3_1')(deconv_cap_2_2)
    up_3 = layers.Concatenate(axis=-2, name='up_3')([deconv_cap_3_1,
        conv1_reshaped])
    seg_caps = ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=16,
        strides=1, padding='same', routings=3, name='seg_caps')(up_3)
    out_seg = Length(num_classes=n_class, seg=True, name='out_seg')(seg_caps)
    _, H, W, C, A = seg_caps.get_shape()
    y = layers.Input(shape=input_shape[:-1] + (1,))
    masked_by_y = Mask()([seg_caps, y])
    masked = Mask()(seg_caps)

    def shared_decoder(mask_layer):
        recon_remove_dim = layers.Reshape((H.value, W.value, A.value))(
            mask_layer)
        recon_1 = layers.Conv2D(filters=64, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='relu', name='recon_1')(
            recon_remove_dim)
        recon_2 = layers.Conv2D(filters=128, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='relu', name='recon_2')(
            recon_1)
        out_recon = layers.Conv2D(filters=1, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='sigmoid', name=
            'out_recon')(recon_2)
        return out_recon
    train_model = models.Model(inputs=[x, y], outputs=[out_seg,
        shared_decoder(masked_by_y)])
    eval_model = models.Model(inputs=x, outputs=[out_seg, shared_decoder(
        masked)])
    noise = layers.Input(shape=(H.value, W.value, C.value, A.value))
    noised_seg_caps = layers.Add()([seg_caps, noise])
    masked_noised_y = Mask()([noised_seg_caps, y])
    manipulate_model = models.Model(inputs=[x, y, noise], outputs=
        shared_decoder(masked_noised_y))
    return train_model


class Length(layers.Layer):

    def __init__(self, num_classes, seg=True, **kwargs):
        super(Length, self).__init__(**kwargs)
        if num_classes == 2:
            self.num_classes = 1
        else:
            self.num_classes = num_classes
        self.seg = seg

    def call(self, inputs, **kwargs):
        if inputs.get_shape().ndims == 5:
            assert inputs.get_shape()[-2
                ].value == 1, 'Error: Must have num_capsules = 1 going into Length'
            inputs = K.squeeze(inputs, axis=-2)
        return K.expand_dims(tf.norm(inputs, axis=-1), axis=-1)

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 5:
            input_shape = input_shape[0:-2] + input_shape[-1:]
        if self.seg:
            return input_shape[:-1] + (self.num_classes,)
        else:
            return input_shape[:-1]

    def get_config(self):
        config = {'num_classes': self.num_classes, 'seg': self.seg}
        base_config = super(Length, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class Mask(layers.Layer):

    def __init__(self, resize_masks=False, **kwargs):
        super(Mask, self).__init__(**kwargs)
        self.resize_masks = resize_masks

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input
        elif inputs.get_shape().ndims == 3:
            x = K.sqrt(K.sum(K.square(inputs), -1))
            mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.
                get_shape().as_list()[1])
            masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
        else:
            masked = inputs
        return masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:
            if len(input_shape[0]) == 3:
                return tuple([None, input_shape[0][1] * input_shape[0][2]])
            else:
                return input_shape[0]
        elif len(input_shape) == 3:
            return tuple([None, input_shape[1] * input_shape[2]])
        else:
            return input_shape

    def get_config(self):
        config = {'resize_masks': self.resize_masks}
        base_config = super(Mask, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1,
        padding='same', routings=3, kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
            self.input_num_atoms, self.num_capsule * self.num_atoms],
            initializer=self.kernel_initializer, name='W')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[0] *
            input_shape[1], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.
            strides), padding=self.padding, data_format='channels_last')
        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()
        votes = K.reshape(conv, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.
            value, 1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(space[i], self.kernel_size,
                padding=self.padding, stride=self.strides, dilation=1)
            new_space.append(new_dim)
        return (input_shape[0],) + tuple(new_space) + (self.num_capsule,
            self.num_atoms)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'strides': self.
            strides, 'padding': self.padding, 'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.
            kernel_initializer)}
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2,
        upsamp_type='deconv', padding='same', routings=3,
        kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms * self.scaling * self.scaling], initializer=self.
                kernel_initializer, name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms], initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.num_capsule * self.num_atoms, self.
                input_num_atoms], initializer=self.kernel_initializer, name='W'
                )
        else:
            raise NotImplementedError(
                'Upsampling must be one of: "deconv", "resize", or "subpix"')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[1] *
            input_shape[0], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling,
                self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1),
                padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(
                1, 1), padding='same', data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]
            out_height = deconv_length(self.input_height, self.scaling,
                self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.
                kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.
                num_capsule * self.num_atoms)
            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W,
                output_shape, (self.scaling, self.scaling), padding=self.
                padding, data_format='channels_last')
        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()
        votes = K.reshape(outputs, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2],
            1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        output_shape[1] = deconv_length(output_shape[1], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms
        return tuple(output_shape)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'scaling': self.
            scaling, 'padding': self.padding, 'upsamp_type': self.
            upsamp_type, 'routings': self.routings, 'kernel_initializer':
            initializers.serialize(self.kernel_initializer)}
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def update_routing(votes, biases, logit_shape, num_dims, input_dim,
    output_dim, num_routing):
    if num_dims == 6:
        votes_t_shape = [5, 0, 1, 2, 3, 4]
        r_t_shape = [1, 2, 3, 4, 5, 0]
    elif num_dims == 4:
        votes_t_shape = [3, 0, 1, 2]
        r_t_shape = [1, 2, 3, 0]
    else:
        raise NotImplementedError('Not implemented')
    votes_trans = tf.transpose(votes, votes_t_shape)
    _, _, _, height, width, caps = votes_trans.get_shape()

    def _body(i, logits, activations):
        """Routing while loop."""
        route = tf.nn.softmax(logits, dim=-1)
        preactivate_unrolled = route * votes_trans
        preact_trans = tf.transpose(preactivate_unrolled, r_t_shape)
        preactivate = tf.reduce_sum(preact_trans, axis=1) + biases
        activation = _squash(preactivate)
        activations = activations.write(i, activation)
        act_3d = K.expand_dims(activation, 1)
        tile_shape = np.ones(num_dims, dtype=np.int32).tolist()
        tile_shape[1] = input_dim
        act_replicated = tf.tile(act_3d, tile_shape)
        distances = tf.reduce_sum(votes * act_replicated, axis=-1)
        logits += distances
        return i + 1, logits, activations
    activations = tf.TensorArray(dtype=tf.float32, size=num_routing,
        clear_after_read=False)
    logits = tf.fill(logit_shape, 0.0)
    i = tf.constant(0, dtype=tf.int32)
    _, logits, activations = tf.while_loop(lambda i, logits, activations: i <
        num_routing, _body, loop_vars=[i, logits, activations], swap_memory
        =True)
    return K.cast(activations.read(num_routing - 1), dtype='float32')


def _squash(input_tensor):
    norm = tf.norm(input_tensor, axis=-1, keep_dims=True)
    norm_squared = norm * norm
    return input_tensor / norm * (norm_squared / (1 + norm_squared))


def get_model():
    model = CapsNetR3((256, 256, 1))
    opt = Adam(lr=3e-05)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics={
        'out_seg': dice_hard})
    model.load_weights('weights/croppedtransferlesionsegcaps.h5')
    return model


def main_predict(model, path, truemaskfile):
    raw = Image.open(path)
    raw = np.array(raw.resize((256, 256))) / 255.0
    raw = raw[:, :, 0]
    raw = np.expand_dims(raw, 2)
    print(raw.shape)
    true_msk = Image.open(truemaskfile.replace('vol', 'lesion'))
    true_msk = np.array(true_msk.resize((256, 256))) / 255.0
    true_msk = np.expand_dims(true_msk, 2)
    print(true_msk.shape)
    pred = model.predict([np.expand_dims(raw, 0), np.expand_dims(true_msk, 0)])
    msk = pred[0].squeeze()
    msk = np.stack((msk,) * 3, axis=-1)
    print(np.unique(msk))
    msk[msk >= 0.5] = 1
    msk[msk < 0.5] = 0
    return msk

<mask token>


def dice_soft(y_true, y_pred, loss_type='sorensen', axis=[1, 2, 3], smooth=
    1e-05, from_logits=False):
    if not from_logits:
        _epsilon = tf.convert_to_tensor(1e-07, y_pred.dtype.base_dtype)
        y_pred = tf.clip_by_value(y_pred, _epsilon, 1 - _epsilon)
        y_pred = tf.log(y_pred / (1 - y_pred))
    inse = tf.reduce_sum(y_pred * y_true, axis=axis)
    if loss_type == 'jaccard':
        l = tf.reduce_sum(y_pred * y_pred, axis=axis)
        r = tf.reduce_sum(y_true * y_true, axis=axis)
    elif loss_type == 'sorensen':
        l = tf.reduce_sum(y_pred, axis=axis)
        r = tf.reduce_sum(y_true, axis=axis)
    else:
        raise Exception('Unknow loss_type')
    dice = (2.0 * inse + smooth) / (l + r + smooth)
    dice = tf.reduce_mean(dice)
    return dice


<mask token>


def dice_loss(y_true, y_pred, from_logits=False):
    return 1 - dice_soft(y_true, y_pred)


def CapsNetR3(input_shape, n_class=2):
    x = layers.Input(shape=input_shape)
    conv1 = layers.Conv2D(filters=16, kernel_size=5, strides=1, padding=
        'same', activation='relu', name='conv1')(x)
    _, H, W, C = conv1.get_shape()
    conv1_reshaped = layers.Reshape((H.value, W.value, 1, C.value))(conv1)
    primary_caps = ConvCapsuleLayer(kernel_size=5, num_capsule=2, num_atoms
        =16, strides=2, padding='same', routings=1, name='primarycaps')(
        conv1_reshaped)
    conv_cap_2_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms
        =16, strides=1, padding='same', routings=3, name='conv_cap_2_1')(
        primary_caps)
    conv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms
        =32, strides=2, padding='same', routings=3, name='conv_cap_2_2')(
        conv_cap_2_1)
    conv_cap_3_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =32, strides=1, padding='same', routings=3, name='conv_cap_3_1')(
        conv_cap_2_2)
    conv_cap_3_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =64, strides=2, padding='same', routings=3, name='conv_cap_3_2')(
        conv_cap_3_1)
    conv_cap_4_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms
        =32, strides=1, padding='same', routings=3, name='conv_cap_4_1')(
        conv_cap_3_2)
    deconv_cap_1_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=8,
        num_atoms=32, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_1_1')(conv_cap_4_1)
    up_1 = layers.Concatenate(axis=-2, name='up_1')([deconv_cap_1_1,
        conv_cap_3_1])
    deconv_cap_1_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4,
        num_atoms=32, strides=1, padding='same', routings=3, name=
        'deconv_cap_1_2')(up_1)
    deconv_cap_2_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=4,
        num_atoms=16, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_2_1')(deconv_cap_1_2)
    up_2 = layers.Concatenate(axis=-2, name='up_2')([deconv_cap_2_1,
        conv_cap_2_1])
    deconv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4,
        num_atoms=16, strides=1, padding='same', routings=3, name=
        'deconv_cap_2_2')(up_2)
    deconv_cap_3_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=2,
        num_atoms=16, upsamp_type='deconv', scaling=2, padding='same',
        routings=3, name='deconv_cap_3_1')(deconv_cap_2_2)
    up_3 = layers.Concatenate(axis=-2, name='up_3')([deconv_cap_3_1,
        conv1_reshaped])
    seg_caps = ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=16,
        strides=1, padding='same', routings=3, name='seg_caps')(up_3)
    out_seg = Length(num_classes=n_class, seg=True, name='out_seg')(seg_caps)
    _, H, W, C, A = seg_caps.get_shape()
    y = layers.Input(shape=input_shape[:-1] + (1,))
    masked_by_y = Mask()([seg_caps, y])
    masked = Mask()(seg_caps)

    def shared_decoder(mask_layer):
        recon_remove_dim = layers.Reshape((H.value, W.value, A.value))(
            mask_layer)
        recon_1 = layers.Conv2D(filters=64, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='relu', name='recon_1')(
            recon_remove_dim)
        recon_2 = layers.Conv2D(filters=128, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='relu', name='recon_2')(
            recon_1)
        out_recon = layers.Conv2D(filters=1, kernel_size=1, padding='same',
            kernel_initializer='he_normal', activation='sigmoid', name=
            'out_recon')(recon_2)
        return out_recon
    train_model = models.Model(inputs=[x, y], outputs=[out_seg,
        shared_decoder(masked_by_y)])
    eval_model = models.Model(inputs=x, outputs=[out_seg, shared_decoder(
        masked)])
    noise = layers.Input(shape=(H.value, W.value, C.value, A.value))
    noised_seg_caps = layers.Add()([seg_caps, noise])
    masked_noised_y = Mask()([noised_seg_caps, y])
    manipulate_model = models.Model(inputs=[x, y, noise], outputs=
        shared_decoder(masked_noised_y))
    return train_model


class Length(layers.Layer):

    def __init__(self, num_classes, seg=True, **kwargs):
        super(Length, self).__init__(**kwargs)
        if num_classes == 2:
            self.num_classes = 1
        else:
            self.num_classes = num_classes
        self.seg = seg

    def call(self, inputs, **kwargs):
        if inputs.get_shape().ndims == 5:
            assert inputs.get_shape()[-2
                ].value == 1, 'Error: Must have num_capsules = 1 going into Length'
            inputs = K.squeeze(inputs, axis=-2)
        return K.expand_dims(tf.norm(inputs, axis=-1), axis=-1)

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 5:
            input_shape = input_shape[0:-2] + input_shape[-1:]
        if self.seg:
            return input_shape[:-1] + (self.num_classes,)
        else:
            return input_shape[:-1]

    def get_config(self):
        config = {'num_classes': self.num_classes, 'seg': self.seg}
        base_config = super(Length, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class Mask(layers.Layer):

    def __init__(self, resize_masks=False, **kwargs):
        super(Mask, self).__init__(**kwargs)
        self.resize_masks = resize_masks

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input
        elif inputs.get_shape().ndims == 3:
            x = K.sqrt(K.sum(K.square(inputs), -1))
            mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.
                get_shape().as_list()[1])
            masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
        else:
            masked = inputs
        return masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:
            if len(input_shape[0]) == 3:
                return tuple([None, input_shape[0][1] * input_shape[0][2]])
            else:
                return input_shape[0]
        elif len(input_shape) == 3:
            return tuple([None, input_shape[1] * input_shape[2]])
        else:
            return input_shape

    def get_config(self):
        config = {'resize_masks': self.resize_masks}
        base_config = super(Mask, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1,
        padding='same', routings=3, kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
            self.input_num_atoms, self.num_capsule * self.num_atoms],
            initializer=self.kernel_initializer, name='W')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[0] *
            input_shape[1], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.
            strides), padding=self.padding, data_format='channels_last')
        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()
        votes = K.reshape(conv, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.
            value, 1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(space[i], self.kernel_size,
                padding=self.padding, stride=self.strides, dilation=1)
            new_space.append(new_dim)
        return (input_shape[0],) + tuple(new_space) + (self.num_capsule,
            self.num_atoms)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'strides': self.
            strides, 'padding': self.padding, 'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.
            kernel_initializer)}
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):

    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2,
        upsamp_type='deconv', padding='same', routings=3,
        kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape
            ) == 5, 'The input Tensor should have shape=[None, input_height, input_width, input_num_capsule, input_num_atoms]'
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms * self.scaling * self.scaling], initializer=self.
                kernel_initializer, name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.input_num_atoms, self.num_capsule * self.
                num_atoms], initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.
                kernel_size, self.num_capsule * self.num_atoms, self.
                input_num_atoms], initializer=self.kernel_initializer, name='W'
                )
        else:
            raise NotImplementedError(
                'Upsampling must be one of: "deconv", "resize", or "subpix"')
        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.
            num_atoms], initializer=initializers.constant(0.1), name='b')
        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [input_shape[1] *
            input_shape[0], self.input_height, self.input_width, self.
            input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.
            input_width, self.input_num_atoms))
        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling,
                self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1),
                padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(
                1, 1), padding='same', data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]
            out_height = deconv_length(self.input_height, self.scaling,
                self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.
                kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.
                num_capsule * self.num_atoms)
            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W,
                output_shape, (self.scaling, self.scaling), padding=self.
                padding, data_format='channels_last')
        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()
        votes = K.reshape(outputs, [input_shape[1], input_shape[0],
            votes_shape[1], votes_shape[2], self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value,
            conv_width.value, self.num_capsule, self.num_atoms))
        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[
            1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2],
            1, 1])
        activations = update_routing(votes=votes, biases=biases_replicated,
            logit_shape=logit_shape, num_dims=6, input_dim=self.
            input_num_capsule, output_dim=self.num_capsule, num_routing=
            self.routings)
        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        output_shape[1] = deconv_length(output_shape[1], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self
            .kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms
        return tuple(output_shape)

    def get_config(self):
        config = {'kernel_size': self.kernel_size, 'num_capsule': self.
            num_capsule, 'num_atoms': self.num_atoms, 'scaling': self.
            scaling, 'padding': self.padding, 'upsamp_type': self.
            upsamp_type, 'routings': self.routings, 'kernel_initializer':
            initializers.serialize(self.kernel_initializer)}
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def update_routing(votes, biases, logit_shape, num_dims, input_dim,
    output_dim, num_routing):
    if num_dims == 6:
        votes_t_shape = [5, 0, 1, 2, 3, 4]
        r_t_shape = [1, 2, 3, 4, 5, 0]
    elif num_dims == 4:
        votes_t_shape = [3, 0, 1, 2]
        r_t_shape = [1, 2, 3, 0]
    else:
        raise NotImplementedError('Not implemented')
    votes_trans = tf.transpose(votes, votes_t_shape)
    _, _, _, height, width, caps = votes_trans.get_shape()

    def _body(i, logits, activations):
        """Routing while loop."""
        route = tf.nn.softmax(logits, dim=-1)
        preactivate_unrolled = route * votes_trans
        preact_trans = tf.transpose(preactivate_unrolled, r_t_shape)
        preactivate = tf.reduce_sum(preact_trans, axis=1) + biases
        activation = _squash(preactivate)
        activations = activations.write(i, activation)
        act_3d = K.expand_dims(activation, 1)
        tile_shape = np.ones(num_dims, dtype=np.int32).tolist()
        tile_shape[1] = input_dim
        act_replicated = tf.tile(act_3d, tile_shape)
        distances = tf.reduce_sum(votes * act_replicated, axis=-1)
        logits += distances
        return i + 1, logits, activations
    activations = tf.TensorArray(dtype=tf.float32, size=num_routing,
        clear_after_read=False)
    logits = tf.fill(logit_shape, 0.0)
    i = tf.constant(0, dtype=tf.int32)
    _, logits, activations = tf.while_loop(lambda i, logits, activations: i <
        num_routing, _body, loop_vars=[i, logits, activations], swap_memory
        =True)
    return K.cast(activations.read(num_routing - 1), dtype='float32')


def _squash(input_tensor):
    norm = tf.norm(input_tensor, axis=-1, keep_dims=True)
    norm_squared = norm * norm
    return input_tensor / norm * (norm_squared / (1 + norm_squared))


def get_model():
    model = CapsNetR3((256, 256, 1))
    opt = Adam(lr=3e-05)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics={
        'out_seg': dice_hard})
    model.load_weights('weights/croppedtransferlesionsegcaps.h5')
    return model


def main_predict(model, path, truemaskfile):
    raw = Image.open(path)
    raw = np.array(raw.resize((256, 256))) / 255.0
    raw = raw[:, :, 0]
    raw = np.expand_dims(raw, 2)
    print(raw.shape)
    true_msk = Image.open(truemaskfile.replace('vol', 'lesion'))
    true_msk = np.array(true_msk.resize((256, 256))) / 255.0
    true_msk = np.expand_dims(true_msk, 2)
    print(true_msk.shape)
    pred = model.predict([np.expand_dims(raw, 0), np.expand_dims(true_msk, 0)])
    msk = pred[0].squeeze()
    msk = np.stack((msk,) * 3, axis=-1)
    print(np.unique(msk))
    msk[msk >= 0.5] = 1
    msk[msk < 0.5] = 0
    return msk


import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
import keras
from keras import layers, models
from keras import initializers
from keras.models import Model
from keras.layers import Conv2D, MaxPooling2D, Input, Conv2DTranspose, Concatenate, BatchNormalization, UpSampling2D
from keras.utils.conv_utils import conv_output_length, deconv_length
from keras.layers import Dropout, Activation
from keras.optimizers import Adam, SGD
from keras.layers.advanced_activations import LeakyReLU
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
from keras import backend as K
from keras.utils import plot_model
import tensorflow as tf
import glob
import random
#import cv2
from random import shuffle


def dice_soft(y_true, y_pred, loss_type='sorensen', axis=[1, 2, 3], smooth=1e-5, from_logits=False):
    if not from_logits:
        _epsilon = tf.convert_to_tensor(1e-7, y_pred.dtype.base_dtype)
        y_pred = tf.clip_by_value(y_pred, _epsilon, 1 - _epsilon)
        y_pred = tf.log(y_pred / (1 - y_pred))

    inse = tf.reduce_sum(y_pred * y_true, axis=axis)
    if loss_type == 'jaccard':
        l = tf.reduce_sum(y_pred * y_pred, axis=axis)
        r = tf.reduce_sum(y_true * y_true, axis=axis)
    elif loss_type == 'sorensen':
        l = tf.reduce_sum(y_pred, axis=axis)
        r = tf.reduce_sum(y_true, axis=axis)
    else:
        raise Exception("Unknow loss_type")
    dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    dice = tf.reduce_mean(dice)
    return dice


def dice_hard(y_true, y_pred, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    y_pred = tf.cast(y_pred > threshold, dtype=tf.float32)
    y_true = tf.cast(y_true > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(y_pred, y_true), axis=axis)
    l = tf.reduce_sum(y_pred, axis=axis)
    r = tf.reduce_sum(y_true, axis=axis)
    hard_dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    hard_dice = tf.reduce_mean(hard_dice)
    return hard_dice


def dice_loss(y_true, y_pred, from_logits=False):
    return 1 - dice_soft(y_true, y_pred)

def CapsNetR3(input_shape, n_class=2):
    x = layers.Input(shape=input_shape)

    # Layer 1: Just a conventional Conv2D layer
    conv1 = layers.Conv2D(filters=16, kernel_size=5, strides=1, padding='same', activation='relu', name='conv1')(x)

    # Reshape layer to be 1 capsule x [filters] atoms
    _, H, W, C = conv1.get_shape()
    conv1_reshaped = layers.Reshape((H.value, W.value, 1, C.value))(conv1)

    # Layer 1: Primary Capsule: Conv cap with routing 1
    primary_caps = ConvCapsuleLayer(kernel_size=5, num_capsule=2, num_atoms=16, strides=2, padding='same',
                                    routings=1, name='primarycaps')(conv1_reshaped)

    # Layer 2: Convolutional Capsule
    conv_cap_2_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=16, strides=1, padding='same',
                                    routings=3, name='conv_cap_2_1')(primary_caps)

    # Layer 2: Convolutional Capsule
    conv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=32, strides=2, padding='same',
                                    routings=3, name='conv_cap_2_2')(conv_cap_2_1)

    # Layer 3: Convolutional Capsule
    conv_cap_3_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=32, strides=1, padding='same',
                                    routings=3, name='conv_cap_3_1')(conv_cap_2_2)

    # Layer 3: Convolutional Capsule
    conv_cap_3_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=64, strides=2, padding='same',
                                    routings=3, name='conv_cap_3_2')(conv_cap_3_1)

    # Layer 4: Convolutional Capsule
    conv_cap_4_1 = ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=32, strides=1, padding='same',
                                    routings=3, name='conv_cap_4_1')(conv_cap_3_2)

    # Layer 1 Up: Deconvolutional Capsule
    deconv_cap_1_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=8, num_atoms=32, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_1_1')(conv_cap_4_1)

    # Skip connection
    up_1 = layers.Concatenate(axis=-2, name='up_1')([deconv_cap_1_1, conv_cap_3_1])

    # Layer 1 Up: Deconvolutional Capsule
    deconv_cap_1_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=32, strides=1,
                                      padding='same', routings=3, name='deconv_cap_1_2')(up_1)

    # Layer 2 Up: Deconvolutional Capsule
    deconv_cap_2_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=4, num_atoms=16, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_2_1')(deconv_cap_1_2)

    # Skip connection
    up_2 = layers.Concatenate(axis=-2, name='up_2')([deconv_cap_2_1, conv_cap_2_1])

    # Layer 2 Up: Deconvolutional Capsule
    deconv_cap_2_2 = ConvCapsuleLayer(kernel_size=5, num_capsule=4, num_atoms=16, strides=1,
                                      padding='same', routings=3, name='deconv_cap_2_2')(up_2)

    # Layer 3 Up: Deconvolutional Capsule
    deconv_cap_3_1 = DeconvCapsuleLayer(kernel_size=4, num_capsule=2, num_atoms=16, upsamp_type='deconv',
                                        scaling=2, padding='same', routings=3,
                                        name='deconv_cap_3_1')(deconv_cap_2_2)

    # Skip connection
    up_3 = layers.Concatenate(axis=-2, name='up_3')([deconv_cap_3_1, conv1_reshaped])

    # Layer 4: Convolutional Capsule: 1x1
    seg_caps = ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=16, strides=1, padding='same',
                                routings=3, name='seg_caps')(up_3)

    # Layer 4: This is an auxiliary layer to replace each capsule with its length. Just to match the true label's shape.
    out_seg = Length(num_classes=n_class, seg=True, name='out_seg')(seg_caps)

    # Decoder network.
    _, H, W, C, A = seg_caps.get_shape()
    y = layers.Input(shape=input_shape[:-1] + (1,))
    masked_by_y = Mask()([seg_caps, y])  # The true label is used to mask the output of capsule layer. For training
    masked = Mask()(seg_caps)  # Mask using the capsule with maximal length. For prediction

    def shared_decoder(mask_layer):
        recon_remove_dim = layers.Reshape((H.value, W.value, A.value))(mask_layer)

        recon_1 = layers.Conv2D(filters=64, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                activation='relu', name='recon_1')(recon_remove_dim)

        recon_2 = layers.Conv2D(filters=128, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                activation='relu', name='recon_2')(recon_1)

        out_recon = layers.Conv2D(filters=1, kernel_size=1, padding='same', kernel_initializer='he_normal',
                                  activation='sigmoid', name='out_recon')(recon_2)

        return out_recon

    # Models for training and evaluation (prediction)
    train_model = models.Model(inputs=[x, y], outputs=[out_seg, shared_decoder(masked_by_y)])
    eval_model = models.Model(inputs=x, outputs=[out_seg, shared_decoder(masked)])

    # manipulate model
    noise = layers.Input(shape=((H.value, W.value, C.value, A.value)))
    noised_seg_caps = layers.Add()([seg_caps, noise])
    masked_noised_y = Mask()([noised_seg_caps, y])
    manipulate_model = models.Model(inputs=[x, y, noise], outputs=shared_decoder(masked_noised_y))

    # return train_model, eval_model, manipulate_model

    return train_model


class Length(layers.Layer):
    def __init__(self, num_classes, seg=True, **kwargs):
        super(Length, self).__init__(**kwargs)
        if num_classes == 2:
            self.num_classes = 1
        else:
            self.num_classes = num_classes
        self.seg = seg

    def call(self, inputs, **kwargs):
        if inputs.get_shape().ndims == 5:
            assert inputs.get_shape()[-2].value == 1, 'Error: Must have num_capsules = 1 going into Length'
            inputs = K.squeeze(inputs, axis=-2)
        return K.expand_dims(tf.norm(inputs, axis=-1), axis=-1)

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 5:
            input_shape = input_shape[0:-2] + input_shape[-1:]
        if self.seg:
            return input_shape[:-1] + (self.num_classes,)
        else:
            return input_shape[:-1]

    def get_config(self):
        config = {'num_classes': self.num_classes, 'seg': self.seg}
        base_config = super(Length, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class Mask(layers.Layer):
    def __init__(self, resize_masks=False, **kwargs):
        super(Mask, self).__init__(**kwargs)
        self.resize_masks = resize_masks

    def call(self, inputs, **kwargs):
        if type(inputs) is list:
            assert len(inputs) == 2
            input, mask = inputs
            _, hei, wid, _, _ = input.get_shape()
            if self.resize_masks:
                mask = tf.image.resize_bicubic(mask, (hei.value, wid.value))
            mask = K.expand_dims(mask, -1)
            if input.get_shape().ndims == 3:
                masked = K.batch_flatten(mask * input)
            else:
                masked = mask * input

        else:
            if inputs.get_shape().ndims == 3:
                x = K.sqrt(K.sum(K.square(inputs), -1))
                mask = K.one_hot(indices=K.argmax(x, 1), num_classes=x.get_shape().as_list()[1])
                masked = K.batch_flatten(K.expand_dims(mask, -1) * inputs)
            else:
                masked = inputs

        return masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:  # true label provided
            if len(input_shape[0]) == 3:
                return tuple([None, input_shape[0][1] * input_shape[0][2]])
            else:
                return input_shape[0]
        else:  # no true label provided
            if len(input_shape) == 3:
                return tuple([None, input_shape[1] * input_shape[2]])
            else:
                return input_shape

    def get_config(self):
        config = {'resize_masks': self.resize_masks}
        base_config = super(Mask, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConvCapsuleLayer(layers.Layer):
    def __init__(self, kernel_size, num_capsule, num_atoms, strides=1, padding='same', routings=3,
                 kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 5, "The input Tensor should have shape=[None, input_height, input_width," \
                                      " input_num_capsule, input_num_atoms]"
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]

        # Transform matrix
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                        self.input_num_atoms, self.num_capsule * self.num_atoms],
                                 initializer=self.kernel_initializer,
                                 name='W')

        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.num_atoms],
                                 initializer=initializers.constant(0.1),
                                 name='b')

        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [
            input_shape[0] * input_shape[1], self.input_height, self.input_width, self.input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.input_width, self.input_num_atoms))

        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.strides),
                        padding=self.padding, data_format='channels_last')

        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()

        votes = K.reshape(conv, [input_shape[1], input_shape[0], votes_shape[1], votes_shape[2],
                                 self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value, conv_width.value,
                         self.num_capsule, self.num_atoms))

        logit_shape = K.stack([
            input_shape[1], input_shape[0], votes_shape[1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.value, 1, 1])

        activations = update_routing(
            votes=votes,
            biases=biases_replicated,
            logit_shape=logit_shape,
            num_dims=6,
            input_dim=self.input_num_capsule,
            output_dim=self.num_capsule,
            num_routing=self.routings)

        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(
                space[i],
                self.kernel_size,
                padding=self.padding,
                stride=self.strides,
                dilation=1)
            new_space.append(new_dim)

        return (input_shape[0],) + tuple(new_space) + (self.num_capsule, self.num_atoms)

    def get_config(self):
        config = {
            'kernel_size': self.kernel_size,
            'num_capsule': self.num_capsule,
            'num_atoms': self.num_atoms,
            'strides': self.strides,
            'padding': self.padding,
            'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.kernel_initializer)
        }
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class DeconvCapsuleLayer(layers.Layer):
    def __init__(self, kernel_size, num_capsule, num_atoms, scaling=2, upsamp_type='deconv', padding='same', routings=3,
                 kernel_initializer='he_normal', **kwargs):
        super(DeconvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_capsule = num_capsule
        self.num_atoms = num_atoms
        self.scaling = scaling
        self.upsamp_type = upsamp_type
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 5, "The input Tensor should have shape=[None, input_height, input_width," \
                                      " input_num_capsule, input_num_atoms]"
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.input_num_capsule = input_shape[3]
        self.input_num_atoms = input_shape[4]

        # Transform matrix
        if self.upsamp_type == 'subpix':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.input_num_atoms,
                                            self.num_capsule * self.num_atoms * self.scaling * self.scaling],
                                     initializer=self.kernel_initializer,
                                     name='W')
        elif self.upsamp_type == 'resize':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.input_num_atoms, self.num_capsule * self.num_atoms],
                                     initializer=self.kernel_initializer, name='W')
        elif self.upsamp_type == 'deconv':
            self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                            self.num_capsule * self.num_atoms, self.input_num_atoms],
                                     initializer=self.kernel_initializer, name='W')
        else:
            raise NotImplementedError('Upsampling must be one of: "deconv", "resize", or "subpix"')

        self.b = self.add_weight(shape=[1, 1, self.num_capsule, self.num_atoms],
                                 initializer=initializers.constant(0.1),
                                 name='b')

        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [
            input_shape[1] * input_shape[0], self.input_height, self.input_width, self.input_num_atoms])
        input_tensor_reshaped.set_shape((None, self.input_height, self.input_width, self.input_num_atoms))

        if self.upsamp_type == 'resize':
            upsamp = K.resize_images(input_tensor_reshaped, self.scaling, self.scaling, 'channels_last')
            outputs = K.conv2d(upsamp, kernel=self.W, strides=(1, 1), padding=self.padding, data_format='channels_last')
        elif self.upsamp_type == 'subpix':
            conv = K.conv2d(input_tensor_reshaped, kernel=self.W, strides=(1, 1), padding='same',
                            data_format='channels_last')
            outputs = tf.depth_to_space(conv, self.scaling)
        else:
            batch_size = input_shape[1] * input_shape[0]

            # Infer the dynamic output shape:
            out_height = deconv_length(self.input_height, self.scaling, self.kernel_size, self.padding, None)
            out_width = deconv_length(self.input_width, self.scaling, self.kernel_size, self.padding, None)
            output_shape = (batch_size, out_height, out_width, self.num_capsule * self.num_atoms)

            outputs = K.conv2d_transpose(input_tensor_reshaped, self.W, output_shape, (self.scaling, self.scaling),
                                         padding=self.padding, data_format='channels_last')

        votes_shape = K.shape(outputs)
        _, conv_height, conv_width, _ = outputs.get_shape()

        votes = K.reshape(outputs, [input_shape[1], input_shape[0], votes_shape[1], votes_shape[2],
                                    self.num_capsule, self.num_atoms])
        votes.set_shape((None, self.input_num_capsule, conv_height.value, conv_width.value,
                         self.num_capsule, self.num_atoms))

        logit_shape = K.stack([
            input_shape[1], input_shape[0], votes_shape[1], votes_shape[2], self.num_capsule])
        biases_replicated = K.tile(self.b, [votes_shape[1], votes_shape[2], 1, 1])

        activations = update_routing(
            votes=votes,
            biases=biases_replicated,
            logit_shape=logit_shape,
            num_dims=6,
            input_dim=self.input_num_capsule,
            output_dim=self.num_capsule,
            num_routing=self.routings)

        return activations

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)

        output_shape[1] = deconv_length(output_shape[1], self.scaling, self.kernel_size, self.padding, None)
        output_shape[2] = deconv_length(output_shape[2], self.scaling, self.kernel_size, self.padding, None)
        output_shape[3] = self.num_capsule
        output_shape[4] = self.num_atoms

        return tuple(output_shape)

    def get_config(self):
        config = {
            'kernel_size': self.kernel_size,
            'num_capsule': self.num_capsule,
            'num_atoms': self.num_atoms,
            'scaling': self.scaling,
            'padding': self.padding,
            'upsamp_type': self.upsamp_type,
            'routings': self.routings,
            'kernel_initializer': initializers.serialize(self.kernel_initializer)
        }
        base_config = super(DeconvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def update_routing(votes, biases, logit_shape, num_dims, input_dim, output_dim,
                   num_routing):
    if num_dims == 6:
        votes_t_shape = [5, 0, 1, 2, 3, 4]
        r_t_shape = [1, 2, 3, 4, 5, 0]
    elif num_dims == 4:
        votes_t_shape = [3, 0, 1, 2]
        r_t_shape = [1, 2, 3, 0]
    else:
        raise NotImplementedError('Not implemented')

    votes_trans = tf.transpose(votes, votes_t_shape)
    _, _, _, height, width, caps = votes_trans.get_shape()

    def _body(i, logits, activations):
        """Routing while loop."""
        # route: [batch, input_dim, output_dim, ...]
        route = tf.nn.softmax(logits, dim=-1)
        preactivate_unrolled = route * votes_trans
        preact_trans = tf.transpose(preactivate_unrolled, r_t_shape)
        preactivate = tf.reduce_sum(preact_trans, axis=1) + biases
        activation = _squash(preactivate)
        activations = activations.write(i, activation)
        act_3d = K.expand_dims(activation, 1)
        tile_shape = np.ones(num_dims, dtype=np.int32).tolist()
        tile_shape[1] = input_dim
        act_replicated = tf.tile(act_3d, tile_shape)
        distances = tf.reduce_sum(votes * act_replicated, axis=-1)
        logits += distances
        return (i + 1, logits, activations)

    activations = tf.TensorArray(
        dtype=tf.float32, size=num_routing, clear_after_read=False)
    logits = tf.fill(logit_shape, 0.0)

    i = tf.constant(0, dtype=tf.int32)
    _, logits, activations = tf.while_loop(
        lambda i, logits, activations: i < num_routing,
        _body,
        loop_vars=[i, logits, activations],
        swap_memory=True)

    return K.cast(activations.read(num_routing - 1), dtype='float32')

def _squash(input_tensor):
    norm = tf.norm(input_tensor, axis=-1, keep_dims=True)
    norm_squared = norm * norm
    return (input_tensor / norm) * (norm_squared / (1 + norm_squared))
def get_model():
    model = CapsNetR3((256, 256, 1))
    opt = Adam(lr=3e-5)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics={'out_seg': dice_hard})
    model.load_weights('weights/croppedtransferlesionsegcaps.h5')
    return model

def main_predict(model, path,truemaskfile):
    raw = Image.open(path)
    raw = np.array(raw.resize((256, 256))) / 255.
    raw=raw[:,:,0]
    raw = np.expand_dims(raw, 2)
    print(raw.shape)
    true_msk = Image.open(truemaskfile.replace('vol', 'lesion'))
    true_msk = np.array(true_msk.resize((256, 256))) / 255.
    true_msk = np.expand_dims(true_msk, 2)
    print(true_msk.shape)
    pred = model.predict([np.expand_dims(raw, 0), np.expand_dims(true_msk, 0)])
    msk = pred[0].squeeze()
    msk = np.stack((msk,) * 3, axis=-1)
    print(np.unique(msk))
    msk[msk >= 0.5] = 1
    msk[msk < 0.5] = 0
    return msk

