Update app.py

app.py

@@ -16,9 +16,8 @@ from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
 from tensorflow.keras.preprocessing import image
 from tensorflow.keras.models import Model
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
 
-# Переменные ================================
+os.environ["KERAS_BACKEND"] = "tensorflow"
 start_token = "[BOS]"
 end_token = "[EOS]"
 cls_token = "[CLS]"
@@ -43,17 +42,13 @@ attn_pool_dim = proj_dim
 attn_pool_heads = num_heads
 cap_query_num = 128
 
-#RNN
 rnn_embedding_dim = 256
 rnn_proj_dim = 512
 
-# =================================
 
-# Загрузка word_index
 with open('vocabs/word_index.json', 'r', encoding='utf-8') as f:
     word_index = {np.str_(word): np.int64(idx) for word, idx in json.load(f).items()}
 
-# Загрузка index_word
 with open('vocabs/index_word.json', 'r', encoding='utf-8') as f:
     index_word = {np.int64(idx): np.str_(word) for idx, word in json.load(f).items()}
 
@@ -81,7 +76,7 @@ class PositionalEmbedding(layers.Layer):
     return output
 
 
-class AttentionalPooling(tf.keras.layers.Layer):
+class AttentionalPooling(layers.Layer):
     def __init__(self, embed_dim, num_heads=6):
         super().__init__()
         self.embed_dim = embed_dim
@@ -100,7 +95,7 @@ class AttentionalPooling(tf.keras.layers.Layer):
         return self.norm(attn_output)
 
 
-class TransformerBlock(tf.keras.layers.Layer):
+class TransformerBlock(layers.Layer):
   def __init__(self, embed_dim, dense_dim, num_heads, dropout_rate=0.1, ln_epsilon=1e-6, is_multimodal=False, **kwargs):
     super().__init__(**kwargs)
     self.embed_dim = embed_dim
@@ -109,14 +104,12 @@ class TransformerBlock(tf.keras.layers.Layer):
     self.dropout_rate = dropout_rate
     self.ln_epsilon = ln_epsilon
 
-    # Self-Attention
     self.self_attention = layers.MultiHeadAttention(
         num_heads=self.num_heads,
         key_dim=self.embed_dim,
         dropout=self.dropout_rate
     )
 
-    # Cross-Attention
     if is_multimodal:
         self.norm2 = layers.LayerNormalization(epsilon=self.ln_epsilon)
         self.dropout2 = layers.Dropout(self.dropout_rate)
@@ -126,19 +119,15 @@ class TransformerBlock(tf.keras.layers.Layer):
             dropout=self.dropout_rate
         )
 
-
-    # Feed-Forward Network
     self.dense_proj = tf.keras.Sequential([
         layers.Dense(self.dense_dim, activation="gelu"),
         layers.Dropout(self.dropout_rate),
         layers.Dense(self.embed_dim)
     ])
 
-    # Layer Normalization
     self.norm1 = layers.LayerNormalization(epsilon=self.ln_epsilon)
     self.norm3 = layers.LayerNormalization(epsilon=self.ln_epsilon)
 
-    # Dropout
     self.dropout1 = layers.Dropout(self.dropout_rate)
     self.dropout3 = layers.Dropout(self.dropout_rate)
 
@@ -148,11 +137,11 @@ class TransformerBlock(tf.keras.layers.Layer):
     causal_mask = tf.linalg.band_part(tf.ones((seq_len, seq_len), tf.bool), -1, 0)
     return tf.expand_dims(causal_mask, 0)
 
-
+    
   def get_combined_mask(self, causal_mask, padding_mask):
     padding_mask = tf.cast(padding_mask, tf.bool)
 
-    padding_mask = tf.expand_dims(padding_mask, 1)  # (B, 1, L)
+    padding_mask = tf.expand_dims(padding_mask, 1)  
     return causal_mask & padding_mask
 
 
@@ -161,31 +150,28 @@ class TransformerBlock(tf.keras.layers.Layer):
     if mask is not None:
       att_mask = self.get_combined_mask(att_mask, mask)
 
-    # Self-Attention
     x = self.norm1(inputs)
     attention_output_1 = self.self_attention(
         query=x, key=x, value=x, attention_mask=att_mask
     )
     attention_output_1 = self.dropout1(attention_output_1)
-    x = x + attention_output_1  # Add & Norm
-
-    # Cross-Attention
+    x = x + attention_output_1  
+      
     if encoder_outputs is not None:
         x_norm = self.norm2(x)
         attention_output_2 = self.cross_attention(
             query=x_norm, key=encoder_outputs, value=encoder_outputs
         )
         attention_output_2 = self.dropout2(attention_output_2)
-        x = x + attention_output_2  # Add & Norm
+        x = x + attention_output_2  
 
-    # Feed-Forward Network (FFN)
     x_norm = self.norm3(x)
     proj_output = self.dense_proj(x_norm)
     proj_output = self.dropout3(proj_output)
-    return x + proj_output  # Add & Norm
+    return x + proj_output  
 
 
-class UnimodalTextDecoder(tf.keras.layers.Layer):
+class UnimodalTextDecoder(layers.Layer):
     def __init__(self, embed_dim, dense_dim, num_heads, dropout_rate=0.1, ln_epsilon=1e-6, num_layers=4, **kwargs):
         super().__init__()
         self.embed_dim = embed_dim
@@ -201,15 +187,13 @@ class UnimodalTextDecoder(tf.keras.layers.Layer):
         ]
         self.norm = tf.keras.layers.LayerNormalization()
 
-
     def call(self, x, mask=None):
         for layer in self.layers:
             x = layer(inputs=x, mask=mask)
         return self.norm(x)
 
 
-
-class MultimodalTextDecoder(tf.keras.layers.Layer):
+class MultimodalTextDecoder(layers.Layer):
     def __init__(self, embed_dim, dense_dim, num_heads, dropout_rate=0.1, ln_epsilon=1e-6, num_layers=4, **kwargs):
         super().__init__()
         self.embed_dim = embed_dim
@@ -225,7 +209,6 @@ class MultimodalTextDecoder(tf.keras.layers.Layer):
         ]
         self.norm = tf.keras.layers.LayerNormalization()
 
-
     def call(self, x, encoder_outputs, mask=None):
         for layer in self.layers:
             x = layer(inputs=x, encoder_outputs=encoder_outputs, mask=mask)
@@ -302,7 +285,6 @@ class CoCaEncoder(tf.keras.Model):
             name="cap_query"
         )
 
-
     def call(self, input, training=False):
         img_feature = self.vit(input).last_hidden_state
 
@@ -316,7 +298,6 @@ class CoCaEncoder(tf.keras.Model):
         return con_feature, cap_feature
 
 
-
 class CoCaDecoder(tf.keras.Model):
     def __init__(self,
             cls_token_id,
@@ -344,7 +325,6 @@ class CoCaDecoder(tf.keras.Model):
 
         self.norm = layers.LayerNormalization()
 
-
     def call(self, inputs, training=False):
         input_text, cap_feature = inputs
         batch_size = tf.shape(input_text)[0]
@@ -366,15 +346,12 @@ class CoCaDecoder(tf.keras.Model):
         return cls_token_feature, multimodal_logits
 
 
-
-# день 6
 class CoCaModel(tf.keras.Model):
     def __init__(self,
         vit,
         cls_token_id,
         num_heads,
         num_layers):
-
         super().__init__()
 
         self.encoder = CoCaEncoder(vit, name="coca_encoder")
@@ -384,34 +361,28 @@ class CoCaModel(tf.keras.Model):
         self.text_to_latents = EmbedToLatents(proj_dim)
 
         self.pad_id = 0
-        self.temperature = 0.2  # 0.5  #0.9 #1.0
+        self.temperature = 0.07
         self.caption_loss_weight = 1.0
         self.contrastive_loss_weight = 1.0
 
         self.perplexity = Perplexity()
 
-
     def call(self, inputs, training=False):
         image, text = inputs
-
         con_feature, cap_feature = self.encoder(image)
         cls_token_feature, multimodal_logits = self.decoder([text, cap_feature])
-
         return con_feature, cls_token_feature, multimodal_logits
 
-
     def compile(self, optimizer):
         super().compile()
         self.optimizer = optimizer
 
-
     def compute_caption_loss(self, multimodal_out, caption_target):
         caption_loss = tf.keras.losses.sparse_categorical_crossentropy(
                 caption_target, multimodal_out, from_logits=True, ignore_class=self.pad_id)
 
         return tf.reduce_mean(caption_loss)
 
-
     def compute_contrastive_loss(self, con_feature, cls_feature):
         text_embeds = tf.squeeze(cls_feature, axis=1)
         image_embeds = tf.squeeze(con_feature, axis=1)
@@ -419,21 +390,17 @@ class CoCaModel(tf.keras.Model):
         text_latents = self.text_to_latents(text_embeds)
         image_latents = self.img_to_latents(image_embeds)
 
-        # Матрица схожести
-        sim = tf.matmul(text_latents, image_latents, transpose_b=True) / self.temperature   # tf.exp(self.log_temp)
+        sim = tf.matmul(text_latents, image_latents, transpose_b=True) / self.temperature   
 
-        # Метки
         batch_size = tf.shape(sim)[0]
         contrastive_labels = tf.range(batch_size)
 
-        # Вычисление потерь
         loss1 = tf.keras.losses.sparse_categorical_crossentropy(contrastive_labels, sim, from_logits=True)
         loss2 = tf.keras.losses.sparse_categorical_crossentropy(contrastive_labels, tf.transpose(sim), from_logits=True)
         contrastive_loss = tf.reduce_mean((loss1 + loss2) * 0.5)
 
         return contrastive_loss
 
-
     def train_step(self, data):
         (images, caption_input), caption_target = data
 
@@ -457,7 +424,6 @@ class CoCaModel(tf.keras.Model):
             'perplexity': self.perplexity.result()
         }
 
-
     def test_step(self, data):
         (images, caption_input), caption_target = data
 
@@ -477,14 +443,10 @@ class CoCaModel(tf.keras.Model):
             'perplexity': self.perplexity.result()
         }
 
-
     def reset_metrics(self):
         self.perplexity.reset_state()
 
 
-# ===========================================
-# Загрузка весов для коки
-
 coca_model = CoCaModel(vit_tiny_model, cls_token_id=cls_token_id, num_heads=num_heads, num_layers=num_layers)
 
 dummy_features = tf.zeros((1, 3, img_size, img_size), dtype=tf.float32)
@@ -498,22 +460,19 @@ save_dir = "models/"
 model_name = "coca_007"
 coca_model.load_weights(f"{save_dir}/{model_name}.weights.h5")
 
-# ===========================================
-# RNN =======================================
+
 img_embed_dim = 2048
 reg_count = 7 * 7
 
 base_model = ResNet50(weights='imagenet', include_top=False)
 model = Model(inputs=base_model.input, outputs=base_model.output)
 
-
 def preprocess_image(img):
     img = tf.image.resize(img, (img_size, img_size))
     img = tf.convert_to_tensor(img)
     img = preprocess_input(img)
     return np.expand_dims(img, axis=0)
 
-
 def create_features(img):
     img = preprocess_image(img)
     features = model.predict(img, verbose=0)
@@ -539,7 +498,6 @@ class BahdanauAttention(layers.Layer):
         return context, alpha
 
 
-
 class ImageCaptioningModel(tf.keras.Model):
     def __init__(self, vocab_size, max_caption_len, embedding_dim=512, lstm_units=512, dropout_rate=0.5, **kwargs):
         super().__init__(**kwargs)
@@ -562,7 +520,6 @@ class ImageCaptioningModel(tf.keras.Model):
 
         self.concatenate = layers.Concatenate(axis=-1)
 
-
     def call(self, inputs):
         features, captions = inputs
 
@@ -588,7 +545,6 @@ class ImageCaptioningModel(tf.keras.Model):
         return self.fc(outputs)
 
 
-
 rnn_model = ImageCaptioningModel(vocab_size, sentence_length-1, rnn_embedding_dim, rnn_proj_dim)
 image_input = np.random.rand(batch_size, reg_count, img_embed_dim).astype(np.float32)
 text_input = np.random.randint(0, 10000, size=(batch_size, sentence_length))
@@ -605,9 +561,6 @@ model_name = "rnn_att_v4"
 
 rnn_model.load_weights(f"{save_dir}/{model_name}.weights.h5")
 
-# =====================================
-# Методы генерации
-
 beam_width=3
 max_length=sentence_length-1
 temperature=1.0
@@ -631,7 +584,6 @@ def has_repeated_ngrams(seq, n=2):
     return len(ngrams) != len(set(ngrams))
 
 
-# метод с улучшениями для коки
 def generate_caption_coca(image):
     img_processed = load_and_preprocess_image(image)
     _, cap_features = coca_model.encoder.predict(img_processed, verbose=0)
@@ -659,7 +611,6 @@ def generate_caption_coca(image):
                 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
 
@@ -673,7 +624,6 @@ def generate_caption_coca(image):
     return " ".join(index_word[i] for i in best_seq if i not in {word_index[start_token], word_index[end_token]})
 
 
-# метод с улучшениями для rnn
 def generate_caption_rnn(image):
     image_embedding = create_features(image)
     beams = [([word_index[start_token]], 0.0)]
@@ -698,7 +648,6 @@ def generate_caption_rnn(image):
                 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
                 new_beams.append((new_seq, new_log_prob))
@@ -717,25 +666,6 @@ def generate_both(image):
     return f"RNN: {caption1}\n\nCoCa: {caption2}"
 
 
-# interface = gr.Interface(
-#     fn=generate_both,
-#     inputs=gr.Image(type="pil", label="Изображение"),
-#     outputs=gr.Textbox(label="Описания", autoscroll=True, show_copy_button=True),
-#     title="Генератор описаний к изображениям",
-#     allow_flagging="never",
-#     submit_btn="Сгенерировать",
-#     clear_btn="Очистить"
-# )
-
-#------------------------------
-css = """
-#hosted-by-hf {
-    top: unset !important;
-    bottom: 20px !important;
-    right: 20px !important;
-}
-"""
-
 interface = gr.Interface(
     fn=generate_both,
     inputs=gr.Image(type="pil", label="Изображение"),
@@ -750,33 +680,7 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("# 🖼️ Генератор описаний к изображениям")
     interface.render()
 
-# if __name__ == "__main__":
-#     #interface.launch(ssr_mode=False)
-#     demo.launch(ssr_mode=False)
-
-
-# custom_css = """
-# footer {visibility: hidden !important;}
-# .share-button {display: none !important;}
-# #component-1 {margin-top: -1.5rem !important;}  # Уменьшаем отступ сверху
-# """
-
-# interface = gr.Interface(
-#     fn=generate_both,
-#     inputs=gr.Image(type="pil", label="Изображение"),
-#     outputs=gr.Textbox(label="Описания", autoscroll=True, show_copy_button=True),
-#     allow_flagging="never",
-#     submit_btn="Сгенерировать",
-#     clear_btn="Очистить"
-# )
-
-# with gr.Blocks(css=custom_css) as demo:  
-#     gr.Markdown("## 🖼️ Генератор описаний к изображениям")
-#     interface.render()
 
 if __name__ == "__main__":
-    demo.launch(
-        ssr_mode=False,
-        show_api=False  
-    )
+    demo.launch(ssr_mode=False, show_api=False)