Upload 6 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Image[[:space:]]2.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,56 @@
+import gradio as gr
+from model_architecture import ImageCaptionGenerationWithAttention
+from transformers import BartForConditionalGeneration, BartTokenizer, ViTModel, ViTImageProcessor
+import torch
+from PIL import Image
+from dotenv import load_dotenv
+import os
+import traceback
+
+load_dotenv()
+HF_TOKEN = os.getenv('hf_token')
+
+
+class GenerateCaptions:
+    def __init__(self):
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
+        vit_model = ViTModel.from_pretrained(
+            "google/vit-base-patch16-224", token=HF_TOKEN).to(self.device)
+        bart_model = BartForConditionalGeneration.from_pretrained(
+            "facebook/bart-base").to(self.device)
+        self.processor = ViTImageProcessor.from_pretrained(
+            "google/vit-base-patch16-224")
+        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
+        self.model = ImageCaptionGenerationWithAttention(
+            vit_model, bart_model, self.tokenizer)
+        self.model.load_state_dict(torch.load(
+            'image_captioning_model_state_dict.pt', map_location=self.device))
+        self.model.eval()
+
+    def generate_caption(self, frame, max_length=50, num_beams=5):
+        try:
+            image_pixel_values = self.processor(
+                frame, return_tensors="pt").pixel_values
+            generated_caption_ids = self.model.generate(
+                image_pixel_values, max_length, num_beams)
+            return self.tokenizer.decode(generated_caption_ids[0], skip_special_tokens=True)
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+
+
+gc = GenerateCaptions()
+
+demo = gr.Interface(
+    fn=gc.generate_caption,
+    inputs=gr.Image(type='pil'),
+    outputs="text",
+    title="Image Caption with Attention",
+    examples=['Image.jpg', 'Image 2.jpg'],
+    submit_btn='Generate Caption',
+    flagging_mode='never'
+)
+
+
+demo.launch()

image_captioning_model_state_dict.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52b231076ad851a143939d672135b36965c3fec9d9d2531c9bd6417207e5a6e0
+size 905995498

model_architecture.py

@@ -0,0 +1,54 @@
+import torch
+from transformers.modeling_outputs import BaseModelOutput
+import torch.nn as nn
+
+
+class ImageCaptionGenerationWithAttention(nn.Module):
+    def __init__(self, vit_model, bart_model, tokenizer):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.vit = vit_model
+        self.bart = bart_model
+        self.visual_projection = nn.Linear(
+            vit_model.config.hidden_size, bart_model.config.d_model)
+
+    def forward(self, pixel_values, input_ids=None, attention_mask=None, labels=None):
+        vit_outputs = self.vit(pixel_values)
+        if isinstance(vit_outputs, tuple):
+            last_hidden_state = vit_outputs[0]
+        else:
+            last_hidden_state = vit_outputs.last_hidden_state
+
+        visual_features = self.visual_projection(last_hidden_state)
+
+        if input_ids is not None:
+            decoder_outputs = self.bart(
+                labels=input_ids,
+                encoder_outputs=BaseModelOutput(
+                    last_hidden_state=visual_features),
+                return_dict=True
+            )
+            return decoder_outputs
+        else:
+            return visual_features
+
+    def generate(self, pixel_values, max_length=50, num_beams=5, early_stopping=True):
+        self.eval()
+        with torch.no_grad():
+            vit_outputs = self.vit(pixel_values)
+            if isinstance(vit_outputs, tuple):
+                last_hidden_state = vit_outputs[0]
+            else:
+                last_hidden_state = vit_outputs.last_hidden_state
+            visual_features = self.visual_projection(last_hidden_state)
+            generated_ids = self.bart.generate(
+                encoder_outputs=BaseModelOutput(
+                    last_hidden_state=visual_features),
+                max_length=max_length,
+                num_beams=num_beams,
+                early_stopping=early_stopping,
+                decoder_start_token_id=self.tokenizer.bos_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                return_dict_in_generate=False
+            )
+            return generated_ids

requirements.txt

@@ -0,0 +1,4 @@
+torch==2.4.1
+transformers==4.35.2
+gradio==5.0.2
+python-dotenv==1.0.1