Update app.py

app.py

@@ -16,8 +16,9 @@ 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"
+
+# Переменные ================================
 start_token = "[BOS]"
 end_token = "[EOS]"
 cls_token = "[CLS]"
@@ -42,13 +43,17 @@ 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()}
 
@@ -76,7 +81,7 @@ class PositionalEmbedding(layers.Layer):
     return output
 
 
-class AttentionalPooling(layers.Layer):
+class AttentionalPooling(tf.keras.layers.Layer):
     def __init__(self, embed_dim, num_heads=6):
         super().__init__()
         self.embed_dim = embed_dim
@@ -95,7 +100,7 @@ class AttentionalPooling(layers.Layer):
         return self.norm(attn_output)
 
 
-class TransformerBlock(layers.Layer):
+class TransformerBlock(tf.keras.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
@@ -104,12 +109,14 @@ class TransformerBlock(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)
@@ -119,15 +126,19 @@ class TransformerBlock(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)
 
@@ -137,11 +148,11 @@ class TransformerBlock(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)  
+    padding_mask = tf.expand_dims(padding_mask, 1)  # (B, 1, L)
     return causal_mask & padding_mask
 
 
@@ -150,28 +161,31 @@ class TransformerBlock(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  
-      
+    x = x + attention_output_1  # Add & Norm
+
+    # Cross-Attention
     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  
+        x = x + attention_output_2  # Add & Norm
 
+    # 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  
+    return x + proj_output  # Add & Norm
 
 
-class UnimodalTextDecoder(layers.Layer):
+class UnimodalTextDecoder(tf.keras.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
@@ -187,13 +201,15 @@ class UnimodalTextDecoder(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(layers.Layer):
+
+class MultimodalTextDecoder(tf.keras.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
@@ -209,6 +225,7 @@ class MultimodalTextDecoder(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)
@@ -285,6 +302,7 @@ class CoCaEncoder(tf.keras.Model):
             name="cap_query"
         )
 
+
     def call(self, input, training=False):
         img_feature = self.vit(input).last_hidden_state
 
@@ -298,6 +316,7 @@ class CoCaEncoder(tf.keras.Model):
         return con_feature, cap_feature
 
 
+
 class CoCaDecoder(tf.keras.Model):
     def __init__(self,
             cls_token_id,
@@ -325,6 +344,7 @@ 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]
@@ -346,12 +366,15 @@ 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")
@@ -361,28 +384,34 @@ class CoCaModel(tf.keras.Model):
         self.text_to_latents = EmbedToLatents(proj_dim)
 
         self.pad_id = 0
-        self.temperature = 0.07
+        self.temperature = 0.2  # 0.5  #0.9 #1.0
         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)
@@ -390,17 +419,21 @@ 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   
+        # Матрица схожести
+        sim = tf.matmul(text_latents, image_latents, transpose_b=True) / self.temperature   # tf.exp(self.log_temp)
 
+        # Метки
         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
 
@@ -424,6 +457,7 @@ class CoCaModel(tf.keras.Model):
             'perplexity': self.perplexity.result()
         }
 
+
     def test_step(self, data):
         (images, caption_input), caption_target = data
 
@@ -443,10 +477,14 @@ 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)
@@ -460,19 +498,22 @@ 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)
@@ -498,6 +539,7 @@ 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)
@@ -520,6 +562,7 @@ class ImageCaptioningModel(tf.keras.Model):
 
         self.concatenate = layers.Concatenate(axis=-1)
 
+
     def call(self, inputs):
         features, captions = inputs
 
@@ -545,6 +588,7 @@ 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))
@@ -561,6 +605,9 @@ 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
@@ -584,6 +631,7 @@ 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)
@@ -611,6 +659,7 @@ 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
 
@@ -624,6 +673,7 @@ 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)]
@@ -648,6 +698,7 @@ 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))
@@ -666,6 +717,25 @@ 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="Изображение"),
@@ -676,11 +746,36 @@ interface = gr.Interface(
     deep_link=False
 )
 
-with gr.Blocks() as demo:
+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  
+    )