enemy7 commited on
Commit
3f6d039
1 Parent(s): 3c7f662

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +580 -0
main.py CHANGED
@@ -22,6 +22,586 @@ sys.stderr =ttt
22
  import os
23
  os.system("python3 -m http.server 7860 -b 0.0.0.0 &")
24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  import time
26
  print("\n\nMODELLLLLLL save ")
27
  os.system('cp log.txt old.txt')
 
22
  import os
23
  os.system("python3 -m http.server 7860 -b 0.0.0.0 &")
24
 
25
+ # -*- coding: utf-8 -*-
26
+ """image_captioning.ipynb
27
+
28
+ Automatically generated by Colaboratory.
29
+
30
+ Original file is located at
31
+ https://colab.research.google.com/#fileId=https%3A//huggingface.co/keras-io/image-captioning/blob/main/image_captioning.ipynb
32
+
33
+ ##### Copyright 2018 The TensorFlow Authors.
34
+ """
35
+
36
+ #@title Licensed under the Apache License, Version 2.0 (the "License");
37
+ # you may not use this file except in compliance with the License.
38
+ # You may obtain a copy of the License at
39
+ #
40
+ # https://www.apache.org/licenses/LICENSE-2.0
41
+ #
42
+ # Unless required by applicable law or agreed to in writing, software
43
+ # distributed under the License is distributed on an "AS IS" BASIS,
44
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
45
+ # See the License for the specific language governing permissions and
46
+ # limitations under the License.
47
+
48
+
49
+
50
+ import tensorflow as tf
51
+
52
+ # You'll generate plots of attention in order to see which parts of an image
53
+ # your model focuses on during captioning
54
+ import matplotlib.pyplot as plt
55
+
56
+ import collections
57
+ import random
58
+ import numpy as np
59
+ import os
60
+ import time
61
+ import json
62
+ from PIL import Image
63
+
64
+ """## Download and prepare the MS-COCO dataset
65
+
66
+ You will use the [MS-COCO dataset](http://cocodataset.org/#home) to train your model. The dataset contains over 82,000 images, each of which has at least 5 different caption annotations. The code below downloads and extracts the dataset automatically.
67
+
68
+ **Caution: large download ahead**. You'll use the training set, which is a 13GB file.
69
+ """
70
+
71
+ # Download caption annotation files
72
+ annotation_folder = '/annotations/'
73
+ if not os.path.exists(os.path.abspath('.') + annotation_folder):
74
+ annotation_zip = tf.keras.utils.get_file('captions.zip',
75
+ cache_subdir=os.path.abspath('.'),
76
+ origin='http://images.cocodataset.org/annotations/annotations_trainval2017.zip',
77
+ extract=True)
78
+ annotation_file = os.path.dirname(annotation_zip)+'/annotations/captions_train2017.json'
79
+ os.remove(annotation_zip)
80
+
81
+ # Download image files
82
+ image_folder = '/train2017/'
83
+ if not os.path.exists(os.path.abspath('.') + image_folder):
84
+ image_zip = tf.keras.utils.get_file('train2017.zip',
85
+ cache_subdir=os.path.abspath('.'),
86
+ origin='http://images.cocodataset.org/zips/train2017.zip',
87
+ extract=True)
88
+ PATH = os.path.dirname(image_zip) + image_folder
89
+ os.remove(image_zip)
90
+ else:
91
+ PATH = os.path.abspath('.') + image_folder
92
+
93
+ PATH
94
+
95
+ """## Optional: limit the size of the training set
96
+ To speed up training for this tutorial, you'll use a subset of 30,000 captions and their corresponding images to train your model. Choosing to use more data would result in improved captioning quality.
97
+ """
98
+
99
+ with open(annotation_file, 'r') as f:
100
+ annotations = json.load(f)
101
+
102
+ # Group all captions together having the same image ID.
103
+ image_path_to_caption = collections.defaultdict(list)
104
+ for val in annotations['annotations']:
105
+ caption = f"<start> {val['caption']} <end>"
106
+ image_path = PATH + 'COCO_train2014_' + '%012d.jpg' % (val['image_id'])
107
+ image_path_to_caption[image_path].append(caption)
108
+
109
+ image_paths = list(image_path_to_caption.keys())
110
+ random.shuffle(image_paths)
111
+
112
+ # Select the first 6000 image_paths from the shuffled set.
113
+ # Approximately each image id has 5 captions associated with it, so that will
114
+ # lead to 30,000 examples.
115
+ train_image_paths = image_paths[:6000]
116
+ print(len(train_image_paths))
117
+
118
+ train_captions = []
119
+ img_name_vector = []
120
+
121
+ for image_path in train_image_paths:
122
+ caption_list = image_path_to_caption[image_path]
123
+ train_captions.extend(caption_list)
124
+ img_name_vector.extend([image_path] * len(caption_list))
125
+
126
+ print(train_captions[0])
127
+ Image.open(img_name_vector[0])
128
+
129
+ """## Preprocess the images using InceptionV3
130
+ Next, you will use InceptionV3 (which is pretrained on Imagenet) to classify each image. You will extract features from the last convolutional layer.
131
+
132
+ First, you will convert the images into InceptionV3's expected format by:
133
+ * Resizing the image to 299px by 299px
134
+ * [Preprocess the images](https://cloud.google.com/tpu/docs/inception-v3-advanced#preprocessing_stage) using the [preprocess_input](https://www.tensorflow.org/api_docs/python/tf/keras/applications/inception_v3/preprocess_input) method to normalize the image so that it contains pixels in the range of -1 to 1, which matches the format of the images used to train InceptionV3.
135
+ """
136
+
137
+ def load_image(image_path):
138
+ img = tf.io.read_file(image_path)
139
+ img = tf.io.decode_jpeg(img, channels=3)
140
+ img = tf.keras.layers.Resizing(299, 299)(img)
141
+ img = tf.keras.applications.inception_v3.preprocess_input(img)
142
+ return img, image_path
143
+
144
+ """## Initialize InceptionV3 and load the pretrained Imagenet weights
145
+
146
+ Now you'll create a tf.keras model where the output layer is the last convolutional layer in the InceptionV3 architecture. The shape of the output of this layer is ```8x8x2048```. You use the last convolutional layer because you are using attention in this example. You don't perform this initialization during training because it could become a bottleneck.
147
+
148
+ * You forward each image through the network and store the resulting vector in a dictionary (image_name --> feature_vector).
149
+ * After all the images are passed through the network, you save the dictionary to disk.
150
+
151
+ """
152
+
153
+ image_model = tf.keras.applications.InceptionV3(include_top=False,
154
+ weights='imagenet')
155
+ new_input = image_model.input
156
+ hidden_layer = image_model.layers[-1].output
157
+
158
+ image_features_extract_model = tf.keras.Model(new_input, hidden_layer)
159
+
160
+ """## Caching the features extracted from InceptionV3
161
+
162
+ You will pre-process each image with InceptionV3 and cache the output to disk. Caching the output in RAM would be faster but also memory intensive, requiring 8 \* 8 \* 2048 floats per image. At the time of writing, this exceeds the memory limitations of Colab (currently 12GB of memory).
163
+
164
+ Performance could be improved with a more sophisticated caching strategy (for example, by sharding the images to reduce random access disk I/O), but that would require more code.
165
+
166
+ The caching will take about 10 minutes to run in Colab with a GPU. If you'd like to see a progress bar, you can:
167
+
168
+ 1. Install [tqdm](https://github.com/tqdm/tqdm):
169
+
170
+ `!pip install tqdm`
171
+
172
+ 2. Import tqdm:
173
+
174
+ `from tqdm import tqdm`
175
+
176
+ 3. Change the following line:
177
+
178
+ `for img, path in image_dataset:`
179
+
180
+ to:
181
+
182
+ `for img, path in tqdm(image_dataset):`
183
+
184
+ """
185
+
186
+ # Get unique images
187
+ encode_train = sorted(set(img_name_vector))
188
+
189
+ # Feel free to change batch_size according to your system configuration
190
+ image_dataset = tf.data.Dataset.from_tensor_slices(encode_train)
191
+ image_dataset = image_dataset.map(
192
+ load_image, num_parallel_calls=tf.data.AUTOTUNE).batch(16)
193
+
194
+ for img, path in image_dataset:
195
+ batch_features = image_features_extract_model(img)
196
+ batch_features = tf.reshape(batch_features,
197
+ (batch_features.shape[0], -1, batch_features.shape[3]))
198
+
199
+ for bf, p in zip(batch_features, path):
200
+ path_of_feature = p.numpy().decode("utf-8")
201
+ np.save(path_of_feature, bf.numpy())
202
+
203
+ """## Preprocess and tokenize the captions
204
+
205
+ You will transform the text captions into integer sequences using the [TextVectorization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/TextVectorization) layer, with the following steps:
206
+
207
+ * Use [adapt](https://www.tensorflow.org/api_docs/python/tf/keras/layers/TextVectorization#adapt) to iterate over all captions, split the captions into words, and compute a vocabulary of the top 5,000 words (to save memory).
208
+ * Tokenize all captions by mapping each word to it's index in the vocabulary. All output sequences will be padded to length 50.
209
+ * Create word-to-index and index-to-word mappings to display results.
210
+ """
211
+
212
+ caption_dataset = tf.data.Dataset.from_tensor_slices(train_captions)
213
+
214
+ # We will override the default standardization of TextVectorization to preserve
215
+ # "<>" characters, so we preserve the tokens for the <start> and <end>.
216
+ def standardize(inputs):
217
+ inputs = tf.strings.lower(inputs)
218
+ return tf.strings.regex_replace(inputs,
219
+ r"!\"#$%&\(\)\*\+.,-/:;=?@\[\\\]^_`{|}~", "")
220
+
221
+ # Max word count for a caption.
222
+ max_length = 50
223
+ # Use the top 5000 words for a vocabulary.
224
+ vocabulary_size = 5000
225
+ tokenizer = tf.keras.layers.TextVectorization(
226
+ max_tokens=vocabulary_size,
227
+ standardize=standardize,
228
+ output_sequence_length=max_length)
229
+ # Learn the vocabulary from the caption data.
230
+ tokenizer.adapt(caption_dataset)
231
+
232
+ # Create the tokenized vectors
233
+ cap_vector = caption_dataset.map(lambda x: tokenizer(x))
234
+
235
+ # Create mappings for words to indices and indicies to words.
236
+ word_to_index = tf.keras.layers.StringLookup(
237
+ mask_token="",
238
+ vocabulary=tokenizer.get_vocabulary())
239
+ index_to_word = tf.keras.layers.StringLookup(
240
+ mask_token="",
241
+ vocabulary=tokenizer.get_vocabulary(),
242
+ invert=True)
243
+
244
+ """## Split the data into training and testing"""
245
+
246
+ img_to_cap_vector = collections.defaultdict(list)
247
+ for img, cap in zip(img_name_vector, cap_vector):
248
+ img_to_cap_vector[img].append(cap)
249
+
250
+ # Create training and validation sets using an 80-20 split randomly.
251
+ img_keys = list(img_to_cap_vector.keys())
252
+ random.shuffle(img_keys)
253
+
254
+ slice_index = int(len(img_keys)*0.8)
255
+ img_name_train_keys, img_name_val_keys = img_keys[:slice_index], img_keys[slice_index:]
256
+
257
+ img_name_train = []
258
+ cap_train = []
259
+ for imgt in img_name_train_keys:
260
+ capt_len = len(img_to_cap_vector[imgt])
261
+ img_name_train.extend([imgt] * capt_len)
262
+ cap_train.extend(img_to_cap_vector[imgt])
263
+
264
+ img_name_val = []
265
+ cap_val = []
266
+ for imgv in img_name_val_keys:
267
+ capv_len = len(img_to_cap_vector[imgv])
268
+ img_name_val.extend([imgv] * capv_len)
269
+ cap_val.extend(img_to_cap_vector[imgv])
270
+
271
+ len(img_name_train), len(cap_train), len(img_name_val), len(cap_val)
272
+
273
+ """## Create a tf.data dataset for training
274
+
275
+ Your images and captions are ready! Next, let's create a `tf.data` dataset to use for training your model.
276
+ """
277
+
278
+ # Feel free to change these parameters according to your system's configuration
279
+
280
+ BATCH_SIZE = 64
281
+ BUFFER_SIZE = 1000
282
+ embedding_dim = 256
283
+ units = 512
284
+ num_steps = len(img_name_train) // BATCH_SIZE
285
+ # Shape of the vector extracted from InceptionV3 is (64, 2048)
286
+ # These two variables represent that vector shape
287
+ features_shape = 2048
288
+ attention_features_shape = 64
289
+
290
+ # Load the numpy files
291
+ def map_func(img_name, cap):
292
+ img_tensor = np.load(img_name.decode('utf-8')+'.npy')
293
+ return img_tensor, cap
294
+
295
+ dataset = tf.data.Dataset.from_tensor_slices((img_name_train, cap_train))
296
+
297
+ # Use map to load the numpy files in parallel
298
+ dataset = dataset.map(lambda item1, item2: tf.numpy_function(
299
+ map_func, [item1, item2], [tf.float32, tf.int64]),
300
+ num_parallel_calls=tf.data.AUTOTUNE)
301
+
302
+ # Shuffle and batch
303
+ dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
304
+ dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
305
+
306
+ """## Model
307
+
308
+ Fun fact: the decoder below is identical to the one in the example for [Neural Machine Translation with Attention](https://www.tensorflow.org/text/tutorials/nmt_with_attention).
309
+
310
+ The model architecture is inspired by the [Show, Attend and Tell](https://arxiv.org/pdf/1502.03044.pdf) paper.
311
+
312
+ * In this example, you extract the features from the lower convolutional layer of InceptionV3 giving us a vector of shape (8, 8, 2048).
313
+ * You squash that to a shape of (64, 2048).
314
+ * This vector is then passed through the CNN Encoder (which consists of a single Fully connected layer).
315
+ * The RNN (here GRU) attends over the image to predict the next word.
316
+ """
317
+
318
+ class BahdanauAttention(tf.keras.Model):
319
+ def __init__(self, units):
320
+ super(BahdanauAttention, self).__init__()
321
+ self.W1 = tf.keras.layers.Dense(units)
322
+ self.W2 = tf.keras.layers.Dense(units)
323
+ self.V = tf.keras.layers.Dense(1)
324
+
325
+ def call(self, features, hidden):
326
+ # features(CNN_encoder output) shape == (batch_size, 64, embedding_dim)
327
+
328
+ # hidden shape == (batch_size, hidden_size)
329
+ # hidden_with_time_axis shape == (batch_size, 1, hidden_size)
330
+ hidden_with_time_axis = tf.expand_dims(hidden, 1)
331
+
332
+ # attention_hidden_layer shape == (batch_size, 64, units)
333
+ attention_hidden_layer = (tf.nn.tanh(self.W1(features) +
334
+ self.W2(hidden_with_time_axis)))
335
+
336
+ # score shape == (batch_size, 64, 1)
337
+ # This gives you an unnormalized score for each image feature.
338
+ score = self.V(attention_hidden_layer)
339
+
340
+ # attention_weights shape == (batch_size, 64, 1)
341
+ attention_weights = tf.nn.softmax(score, axis=1)
342
+
343
+ # context_vector shape after sum == (batch_size, hidden_size)
344
+ context_vector = attention_weights * features
345
+ context_vector = tf.reduce_sum(context_vector, axis=1)
346
+
347
+ return context_vector, attention_weights
348
+
349
+ class CNN_Encoder(tf.keras.Model):
350
+ # Since you have already extracted the features and dumped it
351
+ # This encoder passes those features through a Fully connected layer
352
+ def __init__(self, embedding_dim):
353
+ super(CNN_Encoder, self).__init__()
354
+ # shape after fc == (batch_size, 64, embedding_dim)
355
+ self.fc = tf.keras.layers.Dense(embedding_dim)
356
+
357
+ def call(self, x):
358
+ x = self.fc(x)
359
+ x = tf.nn.relu(x)
360
+ return x
361
+
362
+ class RNN_Decoder(tf.keras.Model):
363
+ def __init__(self, embedding_dim, units, vocab_size):
364
+ super(RNN_Decoder, self).__init__()
365
+ self.units = units
366
+
367
+ self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
368
+ self.gru = tf.keras.layers.GRU(self.units,
369
+ return_sequences=True,
370
+ return_state=True,
371
+ recurrent_initializer='glorot_uniform')
372
+ self.fc1 = tf.keras.layers.Dense(self.units)
373
+ self.fc2 = tf.keras.layers.Dense(vocab_size)
374
+
375
+ self.attention = BahdanauAttention(self.units)
376
+
377
+ def call(self, x, features, hidden):
378
+ # defining attention as a separate model
379
+ context_vector, attention_weights = self.attention(features, hidden)
380
+
381
+ # x shape after passing through embedding == (batch_size, 1, embedding_dim)
382
+ x = self.embedding(x)
383
+
384
+ # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size)
385
+ x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
386
+
387
+ # passing the concatenated vector to the GRU
388
+ output, state = self.gru(x)
389
+
390
+ # shape == (batch_size, max_length, hidden_size)
391
+ x = self.fc1(output)
392
+
393
+ # x shape == (batch_size * max_length, hidden_size)
394
+ x = tf.reshape(x, (-1, x.shape[2]))
395
+
396
+ # output shape == (batch_size * max_length, vocab)
397
+ x = self.fc2(x)
398
+
399
+ return x, state, attention_weights
400
+
401
+ def reset_state(self, batch_size):
402
+ return tf.zeros((batch_size, self.units))
403
+
404
+ encoder = CNN_Encoder(embedding_dim)
405
+ decoder = RNN_Decoder(embedding_dim, units, tokenizer.vocabulary_size())
406
+
407
+ optimizer = tf.keras.optimizers.Adam()
408
+ loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
409
+ from_logits=True, reduction='none')
410
+
411
+
412
+ def loss_function(real, pred):
413
+ mask = tf.math.logical_not(tf.math.equal(real, 0))
414
+ loss_ = loss_object(real, pred)
415
+
416
+ mask = tf.cast(mask, dtype=loss_.dtype)
417
+ loss_ *= mask
418
+
419
+ return tf.reduce_mean(loss_)
420
+
421
+ """## Checkpoint"""
422
+
423
+ checkpoint_path = "./checkpoints/train"
424
+ ckpt = tf.train.Checkpoint(encoder=encoder,
425
+ decoder=decoder,
426
+ optimizer=optimizer)
427
+ ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)
428
+
429
+ start_epoch = 0
430
+ if ckpt_manager.latest_checkpoint:
431
+ start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])
432
+ # restoring the latest checkpoint in checkpoint_path
433
+ ckpt.restore(ckpt_manager.latest_checkpoint)
434
+
435
+ """## Training
436
+
437
+ * You extract the features stored in the respective `.npy` files and then pass those features through the encoder.
438
+ * The encoder output, hidden state(initialized to 0) and the decoder input (which is the start token) is passed to the decoder.
439
+ * The decoder returns the predictions and the decoder hidden state.
440
+ * The decoder hidden state is then passed back into the model and the predictions are used to calculate the loss.
441
+ * Use teacher forcing to decide the next input to the decoder.
442
+ * Teacher forcing is the technique where the target word is passed as the next input to the decoder.
443
+ * The final step is to calculate the gradients and apply it to the optimizer and backpropagate.
444
+
445
+ """
446
+
447
+ # adding this in a separate cell because if you run the training cell
448
+ # many times, the loss_plot array will be reset
449
+ loss_plot = []
450
+
451
+ @tf.function
452
+ def train_step(img_tensor, target):
453
+ loss = 0
454
+
455
+ # initializing the hidden state for each batch
456
+ # because the captions are not related from image to image
457
+ hidden = decoder.reset_state(batch_size=target.shape[0])
458
+
459
+ dec_input = tf.expand_dims([word_to_index('<start>')] * target.shape[0], 1)
460
+
461
+ with tf.GradientTape() as tape:
462
+ features = encoder(img_tensor)
463
+
464
+ for i in range(1, target.shape[1]):
465
+ # passing the features through the decoder
466
+ predictions, hidden, _ = decoder(dec_input, features, hidden)
467
+
468
+ loss += loss_function(target[:, i], predictions)
469
+
470
+ # using teacher forcing
471
+ dec_input = tf.expand_dims(target[:, i], 1)
472
+
473
+ total_loss = (loss / int(target.shape[1]))
474
+
475
+ trainable_variables = encoder.trainable_variables + decoder.trainable_variables
476
+
477
+ gradients = tape.gradient(loss, trainable_variables)
478
+
479
+ optimizer.apply_gradients(zip(gradients, trainable_variables))
480
+
481
+ return loss, total_loss
482
+
483
+ EPOCHS = 50
484
+
485
+ for epoch in range(start_epoch, EPOCHS):
486
+ start = time.time()
487
+ total_loss = 0
488
+
489
+ for (batch, (img_tensor, target)) in enumerate(dataset):
490
+ batch_loss, t_loss = train_step(img_tensor, target)
491
+ total_loss += t_loss
492
+
493
+ if batch % 100 == 0:
494
+ average_batch_loss = batch_loss.numpy()/int(target.shape[1])
495
+ print(f'Epoch {epoch+1} Batch {batch} Loss {average_batch_loss:.4f}')
496
+ # storing the epoch end loss value to plot later
497
+ loss_plot.append(total_loss / num_steps)
498
+
499
+ if epoch % 5 == 0:
500
+ ckpt_manager.save()
501
+
502
+ print(f'Epoch {epoch+1} Loss {total_loss/num_steps:.6f}')
503
+ print(f'Time taken for 1 epoch {time.time()-start:.2f} sec\n')
504
+
505
+ plt.plot(loss_plot)
506
+ plt.xlabel('Epochs')
507
+ plt.ylabel('Loss')
508
+ plt.title('Loss Plot')
509
+ plt.show()
510
+
511
+ """## Caption!
512
+
513
+ * The evaluate function is similar to the training loop, except you don't use teacher forcing here. The input to the decoder at each time step is its previous predictions along with the hidden state and the encoder output.
514
+ * Stop predicting when the model predicts the end token.
515
+ * And store the attention weights for every time step.
516
+ """
517
+
518
+ def evaluate(image):
519
+ attention_plot = np.zeros((max_length, attention_features_shape))
520
+
521
+ hidden = decoder.reset_state(batch_size=1)
522
+
523
+ temp_input = tf.expand_dims(load_image(image)[0], 0)
524
+ img_tensor_val = image_features_extract_model(temp_input)
525
+ img_tensor_val = tf.reshape(img_tensor_val, (img_tensor_val.shape[0],
526
+ -1,
527
+ img_tensor_val.shape[3]))
528
+
529
+ features = encoder(img_tensor_val)
530
+
531
+ dec_input = tf.expand_dims([word_to_index('<start>')], 0)
532
+ result = []
533
+
534
+ for i in range(max_length):
535
+ predictions, hidden, attention_weights = decoder(dec_input,
536
+ features,
537
+ hidden)
538
+
539
+ attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()
540
+
541
+ predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()
542
+ predicted_word = tf.compat.as_text(index_to_word(predicted_id).numpy())
543
+ result.append(predicted_word)
544
+
545
+ if predicted_word == '<end>':
546
+ return result, attention_plot
547
+
548
+ dec_input = tf.expand_dims([predicted_id], 0)
549
+
550
+ attention_plot = attention_plot[:len(result), :]
551
+ return result, attention_plot
552
+
553
+ def plot_attention(image, result, attention_plot):
554
+ temp_image = np.array(Image.open(image))
555
+
556
+ fig = plt.figure(figsize=(10, 10))
557
+
558
+ len_result = len(result)
559
+ for i in range(len_result):
560
+ temp_att = np.resize(attention_plot[i], (8, 8))
561
+ grid_size = max(int(np.ceil(len_result/2)), 2)
562
+ ax = fig.add_subplot(grid_size, grid_size, i+1)
563
+ ax.set_title(result[i])
564
+ img = ax.imshow(temp_image)
565
+ ax.imshow(temp_att, cmap='gray', alpha=0.6, extent=img.get_extent())
566
+
567
+ plt.tight_layout()
568
+ plt.show()
569
+
570
+ # captions on the validation set
571
+ rid = np.random.randint(0, len(img_name_val))
572
+ image = img_name_val[rid]
573
+ real_caption = ' '.join([tf.compat.as_text(index_to_word(i).numpy())
574
+ for i in cap_val[rid] if i not in [0]])
575
+ result, attention_plot = evaluate(image)
576
+
577
+ print('Real Caption:', real_caption)
578
+ print('Prediction Caption:', ' '.join(result))
579
+ plot_attention(image, result, attention_plot)
580
+
581
+ """## Try it on your own images
582
+
583
+ For fun, below you're provided a method you can use to caption your own images with the model you've just trained. Keep in mind, it was trained on a relatively small amount of data, and your images may be different from the training data (so be prepared for weird results!)
584
+
585
+ """
586
+
587
+ image_url = 'https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSsjv8k9FpJH5AvquxbVyd06B5UludsXYeHuTLTGllucw&s'
588
+ image_extension = image_url[-4:]
589
+ image_path = tf.keras.utils.get_file('image'+image_extension, origin=image_url)
590
+
591
+ result, attention_plot = evaluate(image_path)
592
+ print('Prediction Caption:', ' '.join(result))
593
+ # plot_attention(image_path, result, attention_plot)
594
+ # opening the image
595
+ # Image.open(image_path)
596
+
597
+ """# Next steps
598
+
599
+ Congrats! You've just trained an image captioning model with attention. Next, take a look at this example [Neural Machine Translation with Attention](https://www.tensorflow.org/text/tutorials/nmt_with_attention). It uses a similar architecture to translate between Spanish and English sentences. You can also experiment with training the code in this notebook on a different dataset.
600
+ """
601
+
602
+
603
+
604
+
605
  import time
606
  print("\n\nMODELLLLLLL save ")
607
  os.system('cp log.txt old.txt')