Update app.py

app.py

@@ -457,7 +457,7 @@ optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
 coca_model.compile(optimizer)
 
 save_dir = "models/"
-model_name = "coca_v1"
+model_name = "coca"
 coca_model.load_weights(f"{save_dir}/{model_name}.weights.h5")
 
 
@@ -584,44 +584,90 @@ def has_repeated_ngrams(seq, n=2):
     return len(ngrams) != len(set(ngrams))
 
 
+image_mean = [0.5, 0.5, 0.5]
+image_std = [0.5, 0.5, 0.5]
+
+def load_and_preprocess_image(img):
+    #img = tf.image.decode_jpeg(img, channels=3)
+    img = tf.convert_to_tensor(img)
+    img = tf.image.resize(img, (img_size, img_size))
+    img = img / 255.0
+
+    img = (img - image_mean) / image_std
+    img = tf.transpose(img, perm=[2, 0, 1])
+
+    return np.expand_dims(img, axis=0)
+
+
 def generate_caption_coca(image):
     img_processed = load_and_preprocess_image(image)
+    img_processed = np.expand_dims(img_processed, axis=0)
+
     _, cap_features = coca_model.encoder.predict(img_processed, verbose=0)
 
-    beams = [([word_index[start_token]], 0.0)]
+    start_token_id = word_index[start_token]
+    end_token_id = word_index[end_token]
+    sequence = [start_token_id]
+    text_input = np.zeros((1, sentence_length - 1))
 
-    for _ in range(max_length):
-        new_beams = []
-        for seq, log_prob in beams:
-            if seq[-1] == word_index[end_token]:
-                new_beams.append((seq, log_prob))
-                continue
+    for t in range(sentence_length - 1):
+        text_input[0, :len(sequence)] = sequence
 
-            text_input = np.zeros((1, max_length), dtype=np.int32)
-            text_input[0, :len(seq)] = seq
+        _, logits = coca_model.decoder.predict(
+            [text_input, cap_features],
+            verbose=0
+        )
+        next_token = np.argmax(logits[0, t, :])
 
-            predictions = coca_model.decoder.predict([text_input, cap_features], verbose=0)
-            _, logits = predictions  
-            logits = logits[0, len(seq)-1, :] / temperature 
-            probs = np.exp(logits - np.max(logits))
-            probs /= probs.sum()
+        sequence.append(next_token)
+        if next_token == end_token_id or len(sequence) >= (sentence_length - 1):
+            break
 
-            top_k = np.argpartition(probs, -beam_width)[-beam_width:]
-            for token in top_k:
-                new_seq = seq + [token]
-                new_log_prob = (log_prob * len(seq) + np.log(probs[token])) / (len(seq) + 1)
+    caption = " ".join(
+        [index_word[token] for token in sequence
+         if token not in {word_index[start_token], word_index[end_token]}]
+    )
+
+    return caption
+    
+# def generate_caption_coca(image):
+#     img_processed = load_and_preprocess_image(image)
+#     _, cap_features = coca_model.encoder.predict(img_processed, verbose=0)
+
+#     beams = [([word_index[start_token]], 0.0)]
+
+#     for _ in range(max_length):
+#         new_beams = []
+#         for seq, log_prob in beams:
+#             if seq[-1] == word_index[end_token]:
+#                 new_beams.append((seq, log_prob))
+#                 continue
+
+#             text_input = np.zeros((1, max_length), dtype=np.int32)
+#             text_input[0, :len(seq)] = seq
+
+#             predictions = coca_model.decoder.predict([text_input, cap_features], verbose=0)
+#             _, logits = predictions  
+#             logits = logits[0, len(seq)-1, :] / temperature 
+#             probs = np.exp(logits - np.max(logits))
+#             probs /= probs.sum()
+
+#             top_k = np.argpartition(probs, -beam_width)[-beam_width:]
+#             for token in top_k:
+#                 new_seq = seq + [token]
+#                 new_log_prob = (log_prob * len(seq) + np.log(probs[token])) / (len(seq) + 1)
                 
-                if has_repeated_ngrams(new_seq, n=2):
-                    new_log_prob -= 0.5
+#                 if has_repeated_ngrams(new_seq, n=2):
+#                     new_log_prob -= 0.5
 
-                new_beams.append((new_seq, new_log_prob))
+#                 new_beams.append((new_seq, new_log_prob))
 
-        beams = sorted(new_beams, key=lambda x: x[1], reverse=True)[:beam_width]
-        if all(beam[0][-1] == word_index[end_token] for beam in beams):
-            break
+#         beams = sorted(new_beams, key=lambda x: x[1], reverse=True)[:beam_width]
+#         if all(beam[0][-1] == word_index[end_token] for beam in beams):
+#             break
 
-    best_seq = max(beams, key=lambda x: x[1])[0]
-    return " ".join(index_word[i] for i in best_seq if i not in {word_index[start_token], word_index[end_token]})
+#     best_seq = max(beams, key=lambda x: x[1])[0]
+#     return " ".join(index_word[i] for i in best_seq if i not in {word_index[start_token], word_index[end_token]})
 
 
 def generate_caption_rnn(image):