ICM

Sleeping

ICM / model2.py

Update model2.py

fcc4a36 verified almost 3 years ago

16.8 kB

import pickle
import tensorflow as tf
import pandas as pd
import numpy as np
import os

os.environ["KERAS_BACKEND"] = "tensorflow"
import string
import re
import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
import keras
from keras import layers
from keras.applications import efficientnet
from keras.layers import TextVectorization
import nltk
from nltk.translate.bleu_score import corpus_bleu
keras.utils.set_random_seed(111)
VOCAB_SIZE = 8700

# CONTANTS
MAX_LENGTH = 40
IMAGE_SIZE = (299, 299)
SEQ_LENGTH=37
# # Vocabulary size
# VOCAB_SIZE = 10000

# # Fixed length allowed for any sequence
# SEQ_LENGTH = 25

# Dimension for the image embeddings and token embeddings
EMBED_DIM = 512

# Per-layer units in the feed-forward network
FF_DIM = 512

# Other training parameters
BATCH_SIZE = 64 #64
EPOCHS = 50 #30
AUTOTUNE = tf.data.AUTOTUNE
# VOCABULARY_SIZE = 10000
BATCH_SIZE = 32
BUFFER_SIZE = 1000
EMBED_DIM = 512
UNITS = 512


# LOADING DATA
vocab = pickle.load(open('saved_vocabulary/vocab_coco.file', 'rb'))

tokenizer = tf.keras.layers.TextVectorization(
    # max_tokens=VOCABULARY_SIZE,
    standardize=None,
    output_sequence_length=MAX_LENGTH,
    vocabulary=vocab
)

idx2word = tf.keras.layers.StringLookup(
    mask_token="",
    vocabulary=tokenizer.get_vocabulary(),
    invert=True
)


# MODEL
def CNN_Encoder():
    base_model = efficientnet.EfficientNetB0(
        input_shape=(*IMAGE_SIZE, 3),
        include_top=False,
        weights="imagenet",
    )
    # We freeze our feature extractor
    base_model.trainable = False #freezing the weights of layers to prevent them from being updated during training
    base_model_out = base_model.output #retrieves the output tensor from the last layer of EfficientNetB0 model(contains the high-level feature representation extracted from input images)
    base_model_out = layers.Reshape((-1, base_model_out.shape[-1]))(base_model_out) #Reshape the output tensor to have a 2D shape
    cnn_model = keras.models.Model(base_model.input, base_model_out)  # Create the final CNN model using the input layer of the base model and reshaped output
    return cnn_model

class PositionalEmbedding(layers.Layer):
    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
        super().__init__(**kwargs)
        self.token_embeddings = layers.Embedding(
            input_dim=vocab_size, output_dim=embed_dim
        )
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=embed_dim
        )
        self.sequence_length = sequence_length
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))
    def call(self, inputs):
        length = tf.shape(inputs)[-1]
        positions = tf.range(start=0, limit=length, delta=1)
        embedded_tokens = self.token_embeddings(inputs)
        embedded_tokens = embedded_tokens * self.embed_scale
        embedded_positions = self.position_embeddings(positions)
        return embedded_tokens + embedded_positions

    def compute_mask(self, inputs, mask=None):
        return tf.math.not_equal(inputs, 0)
class TransformerEncoderBlock(layers.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim
        self.dense_dim = dense_dim
        self.num_heads = num_heads
        #captures dependencies between different parts of the input sequence.
        self.attention_1 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim, dropout=0.0
        )
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.dense_1 = layers.Dense(embed_dim, activation="relu")

    def call(self, inputs, training, mask=None):
        inputs = self.layernorm_1(inputs)
        inputs = self.dense_1(inputs)

        attention_output_1 = self.attention_1(
            query=inputs,
            value=inputs,
            key=inputs,
            attention_mask=None,
            training=training,
        )
        out_1 = self.layernorm_2(inputs + attention_output_1)
        return out_1

class PositionalEmbedding(layers.Layer):
    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
        super().__init__(**kwargs)
        self.token_embeddings = layers.Embedding(
            input_dim=vocab_size, output_dim=embed_dim
        )
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=embed_dim
        )
        self.sequence_length = sequence_length
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))#4== 2.0

    def call(self, inputs):
        length = tf.shape(inputs)[-1] #gets the length of input sequence or token
        positions = tf.range(start=0, limit=length, delta=1)# Create a tensor of positions from 0 to length-1(if length is 5 then positions = [0,1,2,3,4])
        embedded_tokens = self.token_embeddings(inputs)# Embed tokens using the token embedding layer(self.token_embeddings layer converts these discrete token indices into continuous, high-dimensional vectors (embeddings). Each token in the caption is represented by a dense vector that captures its semantic meaning.)
        embedded_tokens = embedded_tokens * self.embed_scale  # Scale the embedded tokens by the square root of the embedding dimension
        embedded_positions = self.position_embeddings(positions) # Embed positions using the position embedding layer
        return embedded_tokens + embedded_positions 

    def compute_mask(self, inputs, mask=None):
        return tf.math.not_equal(inputs, 0)





class TransformerDecoderBlock(layers.Layer):
    def __init__(self, embed_dim, ff_dim, num_heads, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim
        self.ff_dim = ff_dim
        self.num_heads = num_heads
        #Defines two multi-head self-attention mechanisms (attention_1 and attention_2) with dropout.
        self.attention_1 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim, dropout=0.1
        )
        self.attention_2 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim, dropout=0.1
        )
        #two feed-forward network layers with ReLU activation.
        self.ffn_layer_1 = layers.Dense(ff_dim, activation="relu")
        self.ffn_layer_2 = layers.Dense(embed_dim)
#normalization to the outputs of attention mechanisms 
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.layernorm_3 = layers.LayerNormalization()
#Utilizes a PositionalEmbedding layer to incorporate positional information into the inputs.
        self.embedding = PositionalEmbedding(
            embed_dim=EMBED_DIM,
            sequence_length=SEQ_LENGTH,
            vocab_size=VOCAB_SIZE,
        )
       
   #softmax activation for generating predictions.
        self.out = layers.Dense(VOCAB_SIZE, activation="softmax")

        self.dropout_1 = layers.Dropout(0.3)
        self.dropout_2 = layers.Dropout(0.5)
        self.supports_masking = True # can be useful for handling variable-length sequences.

    def call(self, inputs, encoder_outputs, training, mask=None):
        inputs = self.embedding(inputs)
        causal_mask = self.get_causal_attention_mask(inputs)

        if mask is not None:
            padding_mask = tf.cast(mask[:, :, tf.newaxis], dtype=tf.int32)
            combined_mask = tf.cast(mask[:, tf.newaxis, :], dtype=tf.int32)
            combined_mask = tf.minimum(combined_mask, causal_mask)

        attention_output_1 = self.attention_1(
            query=inputs,
            value=inputs,
            key=inputs,
            attention_mask=combined_mask,
            training=training,
        )
        out_1 = self.layernorm_1(inputs + attention_output_1)

        attention_output_2 = self.attention_2(
            query=out_1,
            value=encoder_outputs,
            key=encoder_outputs,
            attention_mask=padding_mask,
            training=training,
        )
        out_2 = self.layernorm_2(out_1 + attention_output_2)

        ffn_out = self.ffn_layer_1(out_2)
        ffn_out = self.dropout_1(ffn_out, training=training)
        ffn_out = self.ffn_layer_2(ffn_out)

        ffn_out = self.layernorm_3(ffn_out + out_2, training=training)
        ffn_out = self.dropout_2(ffn_out, training=training)
        preds = self.out(ffn_out)
        return preds


    def get_causal_attention_mask(self, inputs):
        input_shape = tf.shape(inputs)
        batch_size, sequence_length = input_shape[0], input_shape[1]
        i = tf.range(sequence_length)[:, tf.newaxis]
        j = tf.range(sequence_length)
        mask = tf.cast(i >= j, dtype="int32")
        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
        mult = tf.concat(
            [
                tf.expand_dims(batch_size, -1),
                tf.constant([1, 1], dtype=tf.int32),
            ],
            axis=0,
        )
        return tf.tile(mask, mult)



class ImageCaptioningModel(keras.Model):
    def __init__(
        self,
        cnn_model,
        encoder,
        decoder,
        num_captions_per_image=5,
        image_aug=None,
    ):
        super().__init__()
        self.cnn_model = cnn_model
        self.encoder = encoder
        self.decoder = decoder
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.acc_tracker = keras.metrics.Mean(name="accuracy")
        self.num_captions_per_image = num_captions_per_image
        self.image_aug = image_aug

    def calculate_loss(self, y_true, y_pred, mask):
        loss = self.loss(y_true, y_pred)
        mask = tf.cast(mask, dtype=loss.dtype)
        loss *= mask
        return tf.reduce_sum(loss) / tf.reduce_sum(mask)

    def calculate_accuracy(self, y_true, y_pred, mask):
        accuracy = tf.equal(y_true, tf.argmax(y_pred, axis=2))
        accuracy = tf.math.logical_and(mask, accuracy)
        accuracy = tf.cast(accuracy, dtype=tf.float32)
        mask = tf.cast(mask, dtype=tf.float32)
        return tf.reduce_sum(accuracy) / tf.reduce_sum(mask)

    def _compute_caption_loss_and_acc(self, img_embed, batch_seq, training=True):
        encoder_out = self.encoder(img_embed, training=training)
        batch_seq_inp = batch_seq[:, :-1]
        batch_seq_true = batch_seq[:, 1:]
        mask = tf.math.not_equal(batch_seq_true, 0)
        batch_seq_pred = self.decoder(
            batch_seq_inp, encoder_out, training=training, mask=mask
        )
        loss = self.calculate_loss(batch_seq_true, batch_seq_pred, mask)
        acc = self.calculate_accuracy(batch_seq_true, batch_seq_pred, mask)
        return loss, acc


    def train_step(self, batch_data):
        batch_img, batch_seq = batch_data
        batch_loss = 0
        batch_acc = 0

        if self.image_aug:
            batch_img = self.image_aug(batch_img)

        # 1. Get image embeddings
        img_embed = self.cnn_model(batch_img)

        # 2. Pass each of the five captions one by one to the decoder
        # along with the encoder outputs and compute the loss as well as accuracy
        # for each caption.
        for i in range(self.num_captions_per_image):
            with tf.GradientTape() as tape:
                loss, acc = self._compute_caption_loss_and_acc(
                    img_embed, batch_seq[:, i, :], training=True
                )

                # 3. Update loss and accuracy
                batch_loss += loss
                batch_acc += acc

            # 4. Get the list of all the trainable weights
            train_vars = (
                self.encoder.trainable_variables + self.decoder.trainable_variables
            )

            # 5. Get the gradients
            grads = tape.gradient(loss, train_vars)

            # 6. Update the trainable weights
            self.optimizer.apply_gradients(zip(grads, train_vars))

        # 7. Update the trackers
        batch_acc /= float(self.num_captions_per_image)
        self.loss_tracker.update_state(batch_loss)
        self.acc_tracker.update_state(batch_acc)

        # 8. Return the loss and accuracy values
        return {
            "loss": self.loss_tracker.result(),
            "acc": self.acc_tracker.result(),
        }

    def test_step(self, batch_data):
        batch_img, batch_seq = batch_data
        batch_loss = 0
        batch_acc = 0

        # 1. Get image embeddings
        img_embed = self.cnn_model(batch_img)

        # 2. Pass each of the five captions one by one to the decoder
        # along with the encoder outputs and compute the loss as well as accuracy
        # for each caption.
        for i in range(self.num_captions_per_image):
            loss, acc = self._compute_caption_loss_and_acc(
                img_embed, batch_seq[:, i, :], training=False
            )

            # 3. Update batch loss and batch accuracy
            batch_loss += loss
            batch_acc += acc

        batch_acc /= float(self.num_captions_per_image)

        # 4. Update the trackers
        self.loss_tracker.update_state(batch_loss)
        self.acc_tracker.update_state(batch_acc)

        # 5. Return the loss and accuracy values
        return {
            "loss": self.loss_tracker.result(),
            "acc": self.acc_tracker.result(),
        }

    @property
    def metrics(self):
        # We need to list our metrics here so the `reset_states()` can be
        # called automatically.
        return [self.loss_tracker, self.acc_tracker]







def load_image_from_path(img_path):
    img = tf.io.read_file(img_path)
    img = tf.io.decode_jpeg(img, channels=3)
    img = tf.keras.layers.Resizing(299, 299)(img)
    img = tf.keras.applications.inception_v3.preprocess_input(img)
    return img


def generate_caption(img, caption_model, add_noise=False):
    if isinstance(img, str):
        img = load_image_from_path(img)
    
    if add_noise == True:
        noise = tf.random.normal(img.shape)*0.1
        img = (img + noise)
        img = (img - tf.reduce_min(img))/(tf.reduce_max(img) - tf.reduce_min(img))
    
    img = tf.expand_dims(img, axis=0)
    img_embed = caption_model.cnn_model(img)
    img_encoded = caption_model.encoder(img_embed, training=False)
    max_decoded_sentence_length = SEQ_LENGTH - 1
    decoded_caption = "<start> "
        for i in range(max_decoded_sentence_length):
            tokenized_caption = vectorization([decoded_caption])[:, :-1]
            mask = tf.math.not_equal(tokenized_caption, 0)
            predictions = caption_model.decoder(
                tokenized_caption, encoded_img, training=False, mask=mask
            )
            sampled_token_index = np.argmax(predictions[0, i, :])
            sampled_token = index_lookup[sampled_token_index]
            if sampled_token == "<end>":
                break
            decoded_caption += " " + sampled_token

        decoded_caption = decoded_caption.replace("<start> ", "")
        decoded_caption = decoded_caption.replace(" <end>", "").strip()

        # Store references and candidates for computing BLEU score
        references = [
            ref.replace("<start> ", "").replace(" <end>", "").strip().split()
            for ref in test_data[sample_img]
        ]
        candidate_tokens = decoded_caption.split()
        all_references.append(references)
        all_candidates.append(candidate_tokens)
    return decoded_caption


def get_caption_model():
    encoder = TransformerEncoderBlock(embed_dim=EMBED_DIM, dense_dim=FF_DIM, num_heads=1)
    decoder = TransformerDecoderBlock(embed_dim=EMBED_DIM, ff_dim=FF_DIM, num_heads=2)
    caption_model = ImageCaptioningModel(cnn_model=CNN_Encoder,
    encoder=encoder,
    decoder=decoder,
    image_aug=None,)
    cnn_model = CNN_Encoder()

    caption_model = ImageCaptioningModel(
        cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=None,
    )

    def call_fn(batch, training):
        return batch

    caption_model.call = call_fn
    sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))

    caption_model((sample_x, sample_y))

    sample_img_embed = caption_model.cnn_model(sample_x)
    sample_enc_out = caption_model.encoder(sample_img_embed, training=False)
    caption_model.decoder(sample_y, sample_enc_out, training=False)

    try:
        caption_model.load_weights('saved_models/image_captioning_coco_weights.h5')
    except FileNotFoundError:
        caption_model.load_weights('Image-Captioning/saved_models/image_captioning_coco_weights.h5')

    return caption_model

	import pickle
	import tensorflow as tf
	import pandas as pd
	import numpy as np
	import os

	os.environ["KERAS_BACKEND"] = "tensorflow"
	import string
	import re
	import numpy as np
	import matplotlib.pyplot as plt

	import tensorflow as tf
	import keras
	from keras import layers
	from keras.applications import efficientnet
	from keras.layers import TextVectorization
	import nltk
	from nltk.translate.bleu_score import corpus_bleu
	keras.utils.set_random_seed(111)
	VOCAB_SIZE = 8700

	# CONTANTS
	MAX_LENGTH = 40
	IMAGE_SIZE = (299, 299)
	SEQ_LENGTH=37
	# # Vocabulary size
	# VOCAB_SIZE = 10000

	# # Fixed length allowed for any sequence
	# SEQ_LENGTH = 25

	# Dimension for the image embeddings and token embeddings
	EMBED_DIM = 512

	# Per-layer units in the feed-forward network
	FF_DIM = 512

	# Other training parameters
	BATCH_SIZE = 64 #64
	EPOCHS = 50 #30
	AUTOTUNE = tf.data.AUTOTUNE
	# VOCABULARY_SIZE = 10000
	BATCH_SIZE = 32
	BUFFER_SIZE = 1000
	EMBED_DIM = 512
	UNITS = 512


	# LOADING DATA
	vocab = pickle.load(open('saved_vocabulary/vocab_coco.file', 'rb'))

	tokenizer = tf.keras.layers.TextVectorization(
	# max_tokens=VOCABULARY_SIZE,
	standardize=None,
	output_sequence_length=MAX_LENGTH,
	vocabulary=vocab
	)

	idx2word = tf.keras.layers.StringLookup(
	mask_token="",
	vocabulary=tokenizer.get_vocabulary(),
	invert=True
	)


	# MODEL
	def CNN_Encoder():
	base_model = efficientnet.EfficientNetB0(
	input_shape=(*IMAGE_SIZE, 3),
	include_top=False,
	weights="imagenet",
	)
	# We freeze our feature extractor
	base_model.trainable = False #freezing the weights of layers to prevent them from being updated during training
	base_model_out = base_model.output #retrieves the output tensor from the last layer of EfficientNetB0 model(contains the high-level feature representation extracted from input images)
	base_model_out = layers.Reshape((-1, base_model_out.shape[-1]))(base_model_out) #Reshape the output tensor to have a 2D shape
	cnn_model = keras.models.Model(base_model.input, base_model_out) # Create the final CNN model using the input layer of the base model and reshaped output
	return cnn_model

	class PositionalEmbedding(layers.Layer):
	def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
	super().__init__(**kwargs)
	self.token_embeddings = layers.Embedding(
	input_dim=vocab_size, output_dim=embed_dim
	)
	self.position_embeddings = layers.Embedding(
	input_dim=sequence_length, output_dim=embed_dim
	)
	self.sequence_length = sequence_length
	self.vocab_size = vocab_size
	self.embed_dim = embed_dim
	self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))
	def call(self, inputs):
	length = tf.shape(inputs)[-1]
	positions = tf.range(start=0, limit=length, delta=1)
	embedded_tokens = self.token_embeddings(inputs)
	embedded_tokens = embedded_tokens * self.embed_scale
	embedded_positions = self.position_embeddings(positions)
	return embedded_tokens + embedded_positions

	def compute_mask(self, inputs, mask=None):
	return tf.math.not_equal(inputs, 0)
	class TransformerEncoderBlock(layers.Layer):
	def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
	super().__init__(**kwargs)
	self.embed_dim = embed_dim
	self.dense_dim = dense_dim
	self.num_heads = num_heads
	#captures dependencies between different parts of the input sequence.
	self.attention_1 = layers.MultiHeadAttention(
	num_heads=num_heads, key_dim=embed_dim, dropout=0.0
	)
	self.layernorm_1 = layers.LayerNormalization()
	self.layernorm_2 = layers.LayerNormalization()
	self.dense_1 = layers.Dense(embed_dim, activation="relu")

	def call(self, inputs, training, mask=None):
	inputs = self.layernorm_1(inputs)
	inputs = self.dense_1(inputs)

	attention_output_1 = self.attention_1(
	query=inputs,
	value=inputs,
	key=inputs,
	attention_mask=None,
	training=training,
	)
	out_1 = self.layernorm_2(inputs + attention_output_1)
	return out_1

	class PositionalEmbedding(layers.Layer):
	def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
	super().__init__(**kwargs)
	self.token_embeddings = layers.Embedding(
	input_dim=vocab_size, output_dim=embed_dim
	)
	self.position_embeddings = layers.Embedding(
	input_dim=sequence_length, output_dim=embed_dim
	)
	self.sequence_length = sequence_length
	self.vocab_size = vocab_size
	self.embed_dim = embed_dim
	self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))#4== 2.0

	def call(self, inputs):
	length = tf.shape(inputs)[-1] #gets the length of input sequence or token
	positions = tf.range(start=0, limit=length, delta=1)# Create a tensor of positions from 0 to length-1(if length is 5 then positions = [0,1,2,3,4])
	embedded_tokens = self.token_embeddings(inputs)# Embed tokens using the token embedding layer(self.token_embeddings layer converts these discrete token indices into continuous, high-dimensional vectors (embeddings). Each token in the caption is represented by a dense vector that captures its semantic meaning.)
	embedded_tokens = embedded_tokens * self.embed_scale # Scale the embedded tokens by the square root of the embedding dimension
	embedded_positions = self.position_embeddings(positions) # Embed positions using the position embedding layer
	return embedded_tokens + embedded_positions

	def compute_mask(self, inputs, mask=None):
	return tf.math.not_equal(inputs, 0)





	class TransformerDecoderBlock(layers.Layer):
	def __init__(self, embed_dim, ff_dim, num_heads, **kwargs):
	super().__init__(**kwargs)
	self.embed_dim = embed_dim
	self.ff_dim = ff_dim
	self.num_heads = num_heads
	#Defines two multi-head self-attention mechanisms (attention_1 and attention_2) with dropout.
	self.attention_1 = layers.MultiHeadAttention(
	num_heads=num_heads, key_dim=embed_dim, dropout=0.1
	)
	self.attention_2 = layers.MultiHeadAttention(
	num_heads=num_heads, key_dim=embed_dim, dropout=0.1
	)
	#two feed-forward network layers with ReLU activation.
	self.ffn_layer_1 = layers.Dense(ff_dim, activation="relu")
	self.ffn_layer_2 = layers.Dense(embed_dim)
	#normalization to the outputs of attention mechanisms
	self.layernorm_1 = layers.LayerNormalization()
	self.layernorm_2 = layers.LayerNormalization()
	self.layernorm_3 = layers.LayerNormalization()
	#Utilizes a PositionalEmbedding layer to incorporate positional information into the inputs.
	self.embedding = PositionalEmbedding(
	embed_dim=EMBED_DIM,
	sequence_length=SEQ_LENGTH,
	vocab_size=VOCAB_SIZE,
	)

	#softmax activation for generating predictions.
	self.out = layers.Dense(VOCAB_SIZE, activation="softmax")

	self.dropout_1 = layers.Dropout(0.3)
	self.dropout_2 = layers.Dropout(0.5)
	self.supports_masking = True # can be useful for handling variable-length sequences.

	def call(self, inputs, encoder_outputs, training, mask=None):
	inputs = self.embedding(inputs)
	causal_mask = self.get_causal_attention_mask(inputs)

	if mask is not None:
	padding_mask = tf.cast(mask[:, :, tf.newaxis], dtype=tf.int32)
	combined_mask = tf.cast(mask[:, tf.newaxis, :], dtype=tf.int32)
	combined_mask = tf.minimum(combined_mask, causal_mask)

	attention_output_1 = self.attention_1(
	query=inputs,
	value=inputs,
	key=inputs,
	attention_mask=combined_mask,
	training=training,
	)
	out_1 = self.layernorm_1(inputs + attention_output_1)

	attention_output_2 = self.attention_2(
	query=out_1,
	value=encoder_outputs,
	key=encoder_outputs,
	attention_mask=padding_mask,
	training=training,
	)
	out_2 = self.layernorm_2(out_1 + attention_output_2)

	ffn_out = self.ffn_layer_1(out_2)
	ffn_out = self.dropout_1(ffn_out, training=training)
	ffn_out = self.ffn_layer_2(ffn_out)

	ffn_out = self.layernorm_3(ffn_out + out_2, training=training)
	ffn_out = self.dropout_2(ffn_out, training=training)
	preds = self.out(ffn_out)
	return preds


	def get_causal_attention_mask(self, inputs):
	input_shape = tf.shape(inputs)
	batch_size, sequence_length = input_shape[0], input_shape[1]
	i = tf.range(sequence_length)[:, tf.newaxis]
	j = tf.range(sequence_length)
	mask = tf.cast(i >= j, dtype="int32")
	mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
	mult = tf.concat(
	[
	tf.expand_dims(batch_size, -1),
	tf.constant([1, 1], dtype=tf.int32),
	],
	axis=0,
	)
	return tf.tile(mask, mult)



	class ImageCaptioningModel(keras.Model):
	def __init__(
	self,
	cnn_model,
	encoder,
	decoder,
	num_captions_per_image=5,
	image_aug=None,
	):
	super().__init__()
	self.cnn_model = cnn_model
	self.encoder = encoder
	self.decoder = decoder
	self.loss_tracker = keras.metrics.Mean(name="loss")
	self.acc_tracker = keras.metrics.Mean(name="accuracy")
	self.num_captions_per_image = num_captions_per_image
	self.image_aug = image_aug

	def calculate_loss(self, y_true, y_pred, mask):
	loss = self.loss(y_true, y_pred)
	mask = tf.cast(mask, dtype=loss.dtype)
	loss *= mask
	return tf.reduce_sum(loss) / tf.reduce_sum(mask)

	def calculate_accuracy(self, y_true, y_pred, mask):
	accuracy = tf.equal(y_true, tf.argmax(y_pred, axis=2))
	accuracy = tf.math.logical_and(mask, accuracy)
	accuracy = tf.cast(accuracy, dtype=tf.float32)
	mask = tf.cast(mask, dtype=tf.float32)
	return tf.reduce_sum(accuracy) / tf.reduce_sum(mask)

	def _compute_caption_loss_and_acc(self, img_embed, batch_seq, training=True):
	encoder_out = self.encoder(img_embed, training=training)
	batch_seq_inp = batch_seq[:, :-1]
	batch_seq_true = batch_seq[:, 1:]
	mask = tf.math.not_equal(batch_seq_true, 0)
	batch_seq_pred = self.decoder(
	batch_seq_inp, encoder_out, training=training, mask=mask
	)
	loss = self.calculate_loss(batch_seq_true, batch_seq_pred, mask)
	acc = self.calculate_accuracy(batch_seq_true, batch_seq_pred, mask)
	return loss, acc


	def train_step(self, batch_data):
	batch_img, batch_seq = batch_data
	batch_loss = 0
	batch_acc = 0

	if self.image_aug:
	batch_img = self.image_aug(batch_img)

	# 1. Get image embeddings
	img_embed = self.cnn_model(batch_img)

	# 2. Pass each of the five captions one by one to the decoder
	# along with the encoder outputs and compute the loss as well as accuracy
	# for each caption.
	for i in range(self.num_captions_per_image):
	with tf.GradientTape() as tape:
	loss, acc = self._compute_caption_loss_and_acc(
	img_embed, batch_seq[:, i, :], training=True
	)

	# 3. Update loss and accuracy
	batch_loss += loss
	batch_acc += acc

	# 4. Get the list of all the trainable weights
	train_vars = (
	self.encoder.trainable_variables + self.decoder.trainable_variables
	)

	# 5. Get the gradients
	grads = tape.gradient(loss, train_vars)

	# 6. Update the trainable weights
	self.optimizer.apply_gradients(zip(grads, train_vars))

	# 7. Update the trackers
	batch_acc /= float(self.num_captions_per_image)
	self.loss_tracker.update_state(batch_loss)
	self.acc_tracker.update_state(batch_acc)

	# 8. Return the loss and accuracy values
	return {
	"loss": self.loss_tracker.result(),
	"acc": self.acc_tracker.result(),
	}

	def test_step(self, batch_data):
	batch_img, batch_seq = batch_data
	batch_loss = 0
	batch_acc = 0

	# 1. Get image embeddings
	img_embed = self.cnn_model(batch_img)

	# 2. Pass each of the five captions one by one to the decoder
	# along with the encoder outputs and compute the loss as well as accuracy
	# for each caption.
	for i in range(self.num_captions_per_image):
	loss, acc = self._compute_caption_loss_and_acc(
	img_embed, batch_seq[:, i, :], training=False
	)

	# 3. Update batch loss and batch accuracy
	batch_loss += loss
	batch_acc += acc

	batch_acc /= float(self.num_captions_per_image)

	# 4. Update the trackers
	self.loss_tracker.update_state(batch_loss)
	self.acc_tracker.update_state(batch_acc)

	# 5. Return the loss and accuracy values
	return {
	"loss": self.loss_tracker.result(),
	"acc": self.acc_tracker.result(),
	}

	@property
	def metrics(self):
	# We need to list our metrics here so the `reset_states()` can be
	# called automatically.
	return [self.loss_tracker, self.acc_tracker]







	def load_image_from_path(img_path):
	img = tf.io.read_file(img_path)
	img = tf.io.decode_jpeg(img, channels=3)
	img = tf.keras.layers.Resizing(299, 299)(img)
	img = tf.keras.applications.inception_v3.preprocess_input(img)
	return img


	def generate_caption(img, caption_model, add_noise=False):
	if isinstance(img, str):
	img = load_image_from_path(img)

	if add_noise == True:
	noise = tf.random.normal(img.shape)*0.1
	img = (img + noise)
	img = (img - tf.reduce_min(img))/(tf.reduce_max(img) - tf.reduce_min(img))

	img = tf.expand_dims(img, axis=0)
	img_embed = caption_model.cnn_model(img)
	img_encoded = caption_model.encoder(img_embed, training=False)
	max_decoded_sentence_length = SEQ_LENGTH - 1
	decoded_caption = "<start> "
	for i in range(max_decoded_sentence_length):
	tokenized_caption = vectorization([decoded_caption])[:, :-1]
	mask = tf.math.not_equal(tokenized_caption, 0)
	predictions = caption_model.decoder(
	tokenized_caption, encoded_img, training=False, mask=mask
	)
	sampled_token_index = np.argmax(predictions[0, i, :])
	sampled_token = index_lookup[sampled_token_index]
	if sampled_token == "<end>":
	break
	decoded_caption += " " + sampled_token

	decoded_caption = decoded_caption.replace("<start> ", "")
	decoded_caption = decoded_caption.replace(" <end>", "").strip()

	# Store references and candidates for computing BLEU score
	references = [
	ref.replace("<start> ", "").replace(" <end>", "").strip().split()
	for ref in test_data[sample_img]
	]
	candidate_tokens = decoded_caption.split()
	all_references.append(references)
	all_candidates.append(candidate_tokens)
	return decoded_caption


	def get_caption_model():
	encoder = TransformerEncoderBlock(embed_dim=EMBED_DIM, dense_dim=FF_DIM, num_heads=1)
	decoder = TransformerDecoderBlock(embed_dim=EMBED_DIM, ff_dim=FF_DIM, num_heads=2)
	caption_model = ImageCaptioningModel(cnn_model=CNN_Encoder,
	encoder=encoder,
	decoder=decoder,
	image_aug=None,)
	cnn_model = CNN_Encoder()

	caption_model = ImageCaptioningModel(
	cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=None,
	)

	def call_fn(batch, training):
	return batch

	caption_model.call = call_fn
	sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))

	caption_model((sample_x, sample_y))

	sample_img_embed = caption_model.cnn_model(sample_x)
	sample_enc_out = caption_model.encoder(sample_img_embed, training=False)
	caption_model.decoder(sample_y, sample_enc_out, training=False)

	try:
	caption_model.load_weights('saved_models/image_captioning_coco_weights.h5')
	except FileNotFoundError:
	caption_model.load_weights('Image-Captioning/saved_models/image_captioning_coco_weights.h5')

	return caption_model