src/load_model.py

import tensorflow as tf
import json
import numpy as np
from tensorflow.keras.models import model_from_json
import os

# Convert the variables to the correct data type
# Load the variables from the JSON file
# Get the directory of the current file (preprocess_coordinates_data.py)
current_directory = os.path.dirname(os.path.abspath(__file__))

# Construct the path to variables.json
json_file_path = os.path.join(current_directory, 'variables.json')
with open(json_file_path, 'r') as json_file:
    variables_dict = json.load(json_file)

# Epsilon value for layer normalisation
LAYER_NORM_EPS = variables_dict['LAYER_NORM_EPS']

# final embedding and transformer embedding size
UNITS_ENCODER = variables_dict['UNITS_ENCODER']
UNITS_DECODER = variables_dict['UNITS_DECODER']

# Transformer
NUM_BLOCKS_ENCODER = variables_dict['NUM_BLOCKS_ENCODER']
NUM_BLOCKS_DECODER = variables_dict['NUM_BLOCKS_DECODER']
NUM_HEADS = variables_dict['NUM_HEADS']
MLP_RATIO = variables_dict['MLP_RATIO']

# Dropout
EMBEDDING_DROPOUT = variables_dict['EMBEDDING_DROPOUT']
MLP_DROPOUT_RATIO = variables_dict['MLP_DROPOUT_RATIO']
MHA_DROPOUT_RATIO = variables_dict['MHA_DROPOUT_RATIO']
CLASSIFIER_DROPOUT_RATIO = variables_dict['CLASSIFIER_DROPOUT_RATIO']

# Number of Frames to resize recording to
N_TARGET_FRAMES = variables_dict['N_TARGET_FRAMES']
N_UNIQUE_CHARACTERS = variables_dict['N_UNIQUE_CHARACTERS']
N_UNIQUE_CHARACTERS0 = variables_dict['N_UNIQUE_CHARACTERS0']
PAD_TOKEN = variables_dict['PAD_TOKEN']
SOS_TOKEN = variables_dict['SOS_TOKEN']

# Length of Phrase + EOS Token
MAX_PHRASE_LENGTH = variables_dict['MAX_PHRASE_LENGTH']

# Mean/Standard Deviations of data used for normalizing
MEANS = np.array(variables_dict['MEANS'])
STDS = np.array(variables_dict['STDS'])

# Initiailizers
INIT_HE_UNIFORM = tf.keras.initializers.he_uniform
INIT_GLOROT_UNIFORM = tf.keras.initializers.glorot_uniform
INIT_ZEROS = tf.keras.initializers.constant(0.0)
# Activations
GELU = tf.keras.activations.gelu


class Embedding(tf.keras.Model):
    def __init__(self, **kwargs):
        super(Embedding, self).__init__(**kwargs)
        self.supports_masking = True

    def build(self, input_shape):
        self.positional_embedding = tf.Variable(
            initial_value=tf.zeros([N_TARGET_FRAMES, UNITS_ENCODER], dtype=tf.float32),
            trainable=True, name='embedding_positional_encoder')
        self.dominant_hand_embedding = LandmarkEmbedding(UNITS_ENCODER, 'dominant_hand')

    def call(self, x, training=False):
        x = tf.where(tf.math.equal(x, 0.0), 0.0, (x - MEANS) / STDS)
        x = self.dominant_hand_embedding(x)
        x = x + self.positional_embedding
        return x

    def get_config(self):
        return super().get_config()

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class Encoder(tf.keras.Model):
    def __init__(self, num_blocks, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        self.num_blocks = num_blocks
        self.support_masking = True
        self.blocks = [
            EncoderTransformerBlock(UNITS_ENCODER, NUM_HEADS, MLP_RATIO, MHA_DROPOUT_RATIO, MLP_DROPOUT_RATIO) for _ in
            range(num_blocks)]

        if UNITS_ENCODER != UNITS_DECODER:
            self.dense_out = tf.keras.layers.Dense(UNITS_DECODER, kernel_initializer=INIT_GLOROT_UNIFORM,
                                                   use_bias=False)
            self.apply_dense_out = True
        else:
            self.apply_dense_out = False

    def call(self, x, x_inp, training=False):
        attention_mask = tf.where(tf.math.reduce_sum(x_inp, axis=[2]) == 0.0, 0.0, 1.0)
        attention_mask = tf.expand_dims(attention_mask, axis=1)
        attention_mask = tf.repeat(attention_mask, repeats=N_TARGET_FRAMES, axis=1)

        for block in self.blocks:
            x = block(x, attention_mask=attention_mask, training=training)

        if self.apply_dense_out:
            x = self.dense_out(x)

        return x, attention_mask

    def get_config(self):
        config = super().get_config()
        config.update({"num_blocks": self.num_blocks})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class Decoder(tf.keras.Model):
    def __init__(self, num_blocks, **kwargs):
        super(Decoder, self).__init__(**kwargs)
        self.num_blocks = num_blocks
        self.supports_masking = True
        self.positional_embedding = tf.Variable(
            initial_value=tf.zeros([N_TARGET_FRAMES, UNITS_DECODER], dtype=tf.float32),
            trainable=True, name='embedding_positional_encoder')
        self.char_emb = tf.keras.layers.Embedding(N_UNIQUE_CHARACTERS, UNITS_DECODER, embeddings_initializer=INIT_ZEROS)
        self.pos_emb_mha = MultiHeadAttention(UNITS_DECODER, NUM_HEADS, MHA_DROPOUT_RATIO)
        self.pos_emb_ln = tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.blocks = [
            DecoderTransformerBlock(UNITS_DECODER, NUM_HEADS, MLP_RATIO, MHA_DROPOUT_RATIO, MLP_DROPOUT_RATIO) for _ in
            range(num_blocks)]

    def get_causal_attention_mask(self, B):
        ones = tf.ones((N_TARGET_FRAMES, N_TARGET_FRAMES))
        mask = tf.linalg.band_part(ones, 0, -1)
        mask = tf.transpose(mask)
        mask = tf.expand_dims(mask, axis=0)
        mask = tf.tile(mask, [B, 1, 1])
        mask = tf.cast(mask, tf.float32)
        return mask

    def call(self, encoder_outputs, attention_mask, phrase, training=False):
        B = tf.shape(encoder_outputs)[0]
        phrase = tf.cast(phrase, tf.int32)
        phrase = tf.pad(phrase, [[0, 0], [1, 0]], constant_values=SOS_TOKEN, name='prepend_sos_token')
        phrase = tf.pad(phrase, [[0, 0], [0, N_TARGET_FRAMES - MAX_PHRASE_LENGTH - 1]], constant_values=PAD_TOKEN,
                        name='append_pad_token')
        causal_mask = self.get_causal_attention_mask(B)
        x = self.positional_embedding + self.char_emb(phrase)
        x = self.pos_emb_ln(x + self.pos_emb_mha(x, x, x, attention_mask=causal_mask))

        for block in self.blocks:
            x = block(x, encoder_outputs, attention_mask=attention_mask, training=training)

        x = tf.slice(x, [0, 0, 0], [-1, MAX_PHRASE_LENGTH, -1])
        return x

    def get_config(self):
        config = super().get_config()
        config.update({"num_blocks": self.num_blocks})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


# Embeds a landmark using fully connected layers
class LandmarkEmbedding(tf.keras.Model):
    def __init__(self, units, name):
        super(LandmarkEmbedding, self).__init__(name=f'{name}_embedding')
        self.units = units
        self.supports_masking = True

    def build(self, input_shape):
        # Embedding for missing landmark in frame, initizlied with zeros
        self.empty_embedding = self.add_weight(
            name=f'{self.name}_empty_embedding',
            shape=[self.units],
            initializer=INIT_ZEROS,
        )
        # Embedding
        self.dense = tf.keras.Sequential([
            tf.keras.layers.Dense(self.units, name=f'{self.name}_dense_1', use_bias=False,
                                  kernel_initializer=INIT_GLOROT_UNIFORM, activation=GELU),  # Can change activation
            tf.keras.layers.Dense(self.units, name=f'{self.name}_dense_2', use_bias=False,
                                  kernel_initializer=INIT_HE_UNIFORM),
        ], name=f'{self.name}_dense')

    def call(self, x):
        return tf.where(
            # Checks whether landmark is missing in frame
            tf.reduce_sum(x, axis=2, keepdims=True) == 0,
            # If so, the empty embedding is used
            self.empty_embedding,
            # Otherwise the landmark data is embedded
            self.dense(x),
        )

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units, "name": self.name})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class EncoderTransformerBlock(tf.keras.layers.Layer):
    def __init__(self, units, num_heads, mlp_ratio, mha_dropout_ratio, mlp_dropout_ratio, **kwargs):
        super(EncoderTransformerBlock, self).__init__(**kwargs)
        self.layer_norm_1 = tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.mha = MultiHeadAttention(units, num_heads, mha_dropout_ratio)
        self.layer_norm_2 = tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.mlp = tf.keras.Sequential([
            tf.keras.layers.Dense(units * mlp_ratio, activation=GELU, kernel_initializer=INIT_GLOROT_UNIFORM,
                                  use_bias=False),
            tf.keras.layers.Dropout(mlp_dropout_ratio),
            tf.keras.layers.Dense(units, kernel_initializer=INIT_HE_UNIFORM, use_bias=False),
        ])

    def call(self, inputs, attention_mask, training=False):
        x = self.layer_norm_1(inputs + self.mha(inputs, inputs, inputs, attention_mask=attention_mask))
        x = self.layer_norm_2(x + self.mlp(x))
        return x

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units, "num_heads": self.num_heads, "mlp_ratio": self.mlp_ratio,
                       "mha_dropout_ratio": self.mha_dropout_ratio, "mlp_dropout_ratio": self.mlp_dropout_ratio})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


# replaced softmax with softmax layer to support masked softmax
def scaled_dot_product(q, k, v, softmax, attention_mask):
    # calculates Q . K(transpose)
    qkt = tf.matmul(q, k, transpose_b=True)
    # calculates scaling factor
    dk = tf.math.sqrt(tf.cast(q.shape[-1], dtype=tf.float32))
    scaled_qkt = qkt / dk
    softmax = softmax(scaled_qkt, mask=attention_mask)
    z = tf.matmul(softmax, v)
    # shape: (m,Tx,depth), same shape as q,k,v
    return z


class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_of_heads, dropout, d_out=None):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_of_heads = num_of_heads
        self.depth = d_model // num_of_heads  # Can change
        self.wq = [tf.keras.layers.Dense(self.depth, use_bias=False) for i in
                   range(num_of_heads)]  # depth//2 isn't common, we can try different numbers
        self.wk = [tf.keras.layers.Dense(self.depth, use_bias=False) for i in range(num_of_heads)]
        self.wv = [tf.keras.layers.Dense(self.depth, use_bias=False) for i in range(num_of_heads)]
        self.softmax = tf.keras.layers.Softmax()
        self.do = tf.keras.layers.Dropout(dropout)
        self.supports_masking = True
        self.wo = tf.keras.layers.Dense(d_model if d_out is None else d_out, use_bias=False)

    def call(self, q, k, v, attention_mask=None, training=False):
        multi_attn = []
        for i in range(self.num_of_heads):
            Q = self.wq[i](q)
            K = self.wk[i](k)
            V = self.wv[i](v)
            multi_attn.append(scaled_dot_product(Q, K, V, self.softmax, attention_mask))

        multi_head = tf.concat(multi_attn, axis=-1)
        multi_head_attention = self.wo(multi_head)
        multi_head_attention = self.do(multi_head_attention, training=training)

        return multi_head_attention

    def get_config(self):
        config = super().get_config()
        config.update({"d_model": self.d_model, "num_of_heads": self.num_of_heads, "dropout": self.dropout})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class DecoderTransformerBlock(tf.keras.layers.Layer):
    def __init__(self, units, num_heads, mlp_ratio, mha_dropout_ratio, mlp_dropout_ratio, **kwargs):
        super(DecoderTransformerBlock, self).__init__(**kwargs)
        self.layer_norm_1 = tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.mha = MultiHeadAttention(units, num_heads, mha_dropout_ratio)
        self.layer_norm_2 = tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.mlp = tf.keras.Sequential([
            tf.keras.layers.Dense(units * mlp_ratio, activation=GELU, kernel_initializer=INIT_GLOROT_UNIFORM,
                                  use_bias=False),
            tf.keras.layers.Dropout(mlp_dropout_ratio),
            tf.keras.layers.Dense(units, kernel_initializer=INIT_HE_UNIFORM, use_bias=False),
        ])

    def call(self, inputs, encoder_outputs, attention_mask, training=False):
        x = self.layer_norm_1(
            inputs + self.mha(inputs, encoder_outputs, encoder_outputs, attention_mask=attention_mask))
        x = self.layer_norm_2(x + self.mlp(x))
        return x

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units, "num_heads": self.num_heads, "mlp_ratio": self.mlp_ratio,
                       "mha_dropout_ratio": self.mha_dropout_ratio, "mlp_dropout_ratio": self.mlp_dropout_ratio})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


custom_objects = {'Embedding': Embedding,
                  'Encoder': Encoder,
                  'Decoder': Decoder,
                  'LandmarkEmbedding': LandmarkEmbedding,
                  'EncoderTransformerBlock': EncoderTransformerBlock,
                  'MultiHeadAttention': MultiHeadAttention,
                  'DecoderTransformerBlock': DecoderTransformerBlock}

# load json and create model
model_archeticture_file_path = os.path.join(current_directory, 'model_architecture.json')
json_file = open(model_archeticture_file_path, 'r')
loaded_model_json = json_file.read()
json_file.close()

# load model from JSON file
loaded_model = model_from_json(loaded_model_json, custom_objects=custom_objects)

# load weights into the new model
model_weights_file_path = os.path.join(current_directory, 'model.h5')
loaded_model.load_weights(model_weights_file_path)

# loaded_model.summary(expand_nested=True, show_trainable=True, )