Files updated

.github/workflows/sync_to_hf.yml

@@ -0,0 +1,21 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab.
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run : |
+          git push --force https://abven:${{ secrets.HF_TOKEN }}@huggingface.co/spaces/abven/ImageCaptionGenerator main

config/config.py

@@ -4,8 +4,8 @@ class Config:
     MAX_SEQ_LEN = 64
     VIT_MODEL = 'google/vit-base-patch16-224-in21k'
     GPT2_MODEL = 'gpt2'
-    LEARNING_RATE = 5e-5
-    EPOCHS = 10
+    LEARNING_RATE = 1e-4 #5e-5
+    EPOCHS = 30
     DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
     AWS_S3_BUCKET = 'your-s3-bucket-name'
     DATASET_PATH = '../Datasets/Flickr8K/'

dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY . .
+
+RUN apt-get update && apt-get install -y build-essential
+
+RUN pip install --upgrade pip
+
+RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Dockerfile
+ENV FLASK_RUN_HOST=0.0.0.0
+
+EXPOSE 5000
+
+CMD ["python", "app.py"]

gradioApp.py

@@ -0,0 +1,37 @@
+import gradio as gr
+from PIL import Image
+import io
+from infer import ImageCaptioningInference
+from models.model import ImageCaptioningModel
+import numpy as np
+
+# Initialize the model
+model_dir = 'model'
+model = ImageCaptioningModel()
+model.load(model_dir)
+inference_model = ImageCaptioningInference(model)
+
+def generate_caption(image):
+    if image is None:
+        return "No image provided."
+    
+    try:
+        # Generate caption using the image path
+        generated_caption = inference_model.infer_image(image)
+        return generated_caption
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Image Captioning",
+    description="Upload an image or select one from your folder to generate a caption.",
+    examples=[["test_img.jpg"]]  # Add some example images if available
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()

models/model.py

@@ -1,7 +1,6 @@
 import torch
 from transformers import ViTModel, ViTFeatureExtractor, GPT2LMHeadModel, GPT2Tokenizer
 from config.config import Config
-from torchsummary import summary
 from torchvision import transforms
 
 class ImageCaptioningModel:
@@ -48,9 +47,8 @@ class ImageCaptioningModel:
     def save(self, path):
         """Save model to disk."""
         self.gpt2_model.save_pretrained(path)
-        self.tokenizer.save_pretrained(path)
 
     def load(self, path):
         """Load model from disk."""
         self.gpt2_model = GPT2LMHeadModel.from_pretrained(path).to(self.device)
-        self.tokenizer = GPT2Tokenizer.from_pretrained(path).to(self.device)
+        # return self.gpt2_model

requirements.txt

@@ -0,0 +1,7 @@
+numpy
+PILLOW
+transformers
+mlflow
+pandas
+torchvision
+flask

train.py

@@ -12,7 +12,6 @@ import mlflow.pytorch
 
 # TODO : Implementing Weights and Biases to for project tracking and evaluation and TODO : DVC also for data versioning
 
-
 def train_model(model,dataLoader, optimizer, loss_fn):
 
     with mlflow.start_run():
@@ -24,7 +23,8 @@ def train_model(model,dataLoader, optimizer, loss_fn):
         })
 
     model.gpt2_model.train()
-    for epoch in range(Config.EPOCHS):
+    for epoch in tqdm(range(Config.EPOCHS)):
+        print(f'Epoch {epoch + 1}/{Config.EPOCHS}')
         epoch_loss = 0
         for batch_idx, (images, captions) in tqdm(enumerate(dataLoader)):
             print(f'\rBatch {batch_idx + 1}/{len(dataLoader)} , Loss : {epoch_loss/(batch_idx+1):.4f}\t', end='')