Update model.py

model.py

@@ -272,7 +272,33 @@ class ImageCaptioningModel(tf.keras.Model):
     @property
     def metrics(self):
         return [self.loss_tracker, self.acc_tracker]
-    
+
+        
+def get_caption_model():
+    encoder = TransformerEncoderLayer(EMBEDDING_DIM, 1)
+    decoder = TransformerDecoderLayer(EMBEDDING_DIM, UNITS, 8)
+
+    cnn_model = CNN_Encoder()
+
+    caption_mode = ImageCaptioningModel(
+        cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=None,
+    )
+
+    def call_fn(batch, training):
+        return batch
+
+    caption_mode.call = call_fn
+    sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))
+
+    caption_mode((sample_x, sample_y))
+
+    sample_img_embed = caption_mode.cnn_model(sample_x)
+    sample_enc_out = caption_mode.encoder(sample_img_embed, training=False)
+    caption_mode.decoder(sample_y, sample_enc_out, training=False)
+
+    caption_mode.load_weights('model.h5')
+
+    return caption_mode
 
 def load_image_from_path(img_path):
     img = tf.io.read_file(img_path)
@@ -312,29 +338,4 @@ def generate_caption(img_path, add_noise=False):
     y_inp = y_inp.replace('[start] ', '')
     return y_inp
 
-def get_caption_model():
-    encoder = TransformerEncoderLayer(EMBEDDING_DIM, 1)
-    decoder = TransformerDecoderLayer(EMBEDDING_DIM, UNITS, 8)
-
-    cnn_model = CNN_Encoder()
-
-    caption_mode = ImageCaptioningModel(
-        cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=None,
-    )
-
-    def call_fn(batch, training):
-        return batch
-
-    caption_mode.call = call_fn
-    sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))
-
-    caption_mode((sample_x, sample_y))
-
-    sample_img_embed = caption_mode.cnn_model(sample_x)
-    sample_enc_out = caption_mode.encoder(sample_img_embed, training=False)
-    caption_mode.decoder(sample_y, sample_enc_out, training=False)
-
-    caption_mode.load_weights('model.h5')
-
-    return caption_mode