Change structure of the project.

app.py

@@ -2,36 +2,33 @@ import torch
 import pandas as pd
 import streamlit as st
 from PIL import Image
-from models.encoder import EncoderCNN
-from models.decoder import DecoderRNN
+
+from encoder import EncoderCNN
+from decoder import DecoderRNN
 from utils.vocab import Vocabulary
-from torchvision import transforms as T
-from utils.helpers import VOCAB_FILE_PATH, CAPTIONS_TKN_PATH
-from inference import sample_with_temp, sample
+#from torchvision import transforms as T
+from utils.helpers import VOCAB_PATH, CAPTIONS_PATH, ENCODER_PATH, DECODER_PATH, IMAGE_DIR
 from utils.transforms import transforms
+from inference import sample_with_temp, sample
 import sacrebleu
-from huggingface_hub import hf_hub_download
 import os
 
 @st.cache_resource
 def load_models():
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-    repo_id = "Sher1988/image-classifier-weights"
-    encoder_path = hf_hub_download(repo_id=repo_id, filename="encoder.pth")
-    decoder_path = hf_hub_download(repo_id=repo_id, filename="decoder.pth")
     
     # Load captions and vocab
-    captions = pd.read_csv(CAPTIONS_TKN_PATH).drop('tokens', axis=1)
-    vocab = Vocabulary(load_path=VOCAB_FILE_PATH)
+    captions = pd.read_csv(CAPTIONS_PATH)
+    vocab = Vocabulary(load_path=VOCAB_PATH)
     
     # Initialize Models
     encoder = EncoderCNN(256).to(device)
     decoder = DecoderRNN(len(vocab), 256, 512).to(device)
     
     # Load Weights
-    encoder.load_state_dict(torch.load(encoder_path, map_location=device))
-    decoder.load_state_dict(torch.load(decoder_path, map_location=device))
+    encoder.load_state_dict(torch.load(ENCODER_PATH, map_location=device))
+    decoder.load_state_dict(torch.load(DECODER_PATH, map_location=device))
     
     encoder.eval()
     decoder.eval()
@@ -41,7 +38,6 @@ def load_models():
 
 # --- Sidebar Configuration ---
 st.sidebar.header("Select an Example Image")
-IMAGE_DIR = "flickr8k/images" # Update this to your local images folder
 
 if os.path.exists(IMAGE_DIR):
     available_images = [f for f in os.listdir(IMAGE_DIR) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
@@ -61,8 +57,8 @@ act_caps = []
 caption = ''
 st.title("📸 AI Image Captioner")
 
-temp = st.slider("Sampling Temperature", min_value=0.1, max_value=2.0, value=0.8, step=0.1)
-st.info("Higher temperature = more creative/random. Lower = more predictable.")
+temp = st.slider("Sampling Temperature", min_value=0.0, max_value=0.8, value=0.1, step=0.1)
+st.info("Higher temperature = more creative/random. Lower temperature = more predictable.")
 
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
@@ -105,7 +101,7 @@ if img is not None:
 
 if act_caps:
     # sacrebleu expects a list of strings for hypothesis 
-    # and a list of lists of strings for references
+    # and a list of strings for references
     refs = [act_caps] 
     sys = [caption]
     

models/decoder.py → decoder.py

@@ -1,24 +1,24 @@
-import torch.nn as nn
-
-
-class DecoderRNN(nn.Module):
-    def __init__(self, vocab_size, embed_size, hidden_size, num_layers=1, padding_idx=0):
-        super(DecoderRNN, self).__init__()
-        self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embed_size, padding_idx=padding_idx)
-        self.lstm = nn.LSTM(input_size=embed_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
-        self.linear = nn.Linear(hidden_size, vocab_size)
-
-        self.init_h = nn.Linear(embed_size, hidden_size)
-        self.init_c = nn.Linear(embed_size, hidden_size)
-    
-    def forward(self, features, captions, hidden=None):
-        if hidden == None:
-            h0 = self.init_h(features).unsqueeze(0).repeat(self.lstm.num_layers, 1, 1)
-            c0 = self.init_c(features).unsqueeze(0).repeat(self.lstm.num_layers, 1, 1)
-            hidden = (h0, c0)
-            # dataflow: (B, seqlen) -> (B, hidden_size) -> (1, B, hidden_size) -> (num_layers, B, hidden_size)
-
-        embeddings = self.embed(captions) # (B, seqlen) ->  Training: (B, seqlen, embed_size) | Inference: (B, 1, embed_size)
-        outputs, hidden = self.lstm(embeddings, hidden) # Training: (B, seqlen, hidden_size) | Inference: (B, 1, hidden_size)
-        outputs = self.linear(outputs) # Training: (B, seqlen, vocab_size) | Inference: (B, 1, vocab_size)
+import torch.nn as nn
+
+
+class DecoderRNN(nn.Module):
+    def __init__(self, vocab_size, embed_size, hidden_size, num_layers=1, padding_idx=0):
+        super(DecoderRNN, self).__init__()
+        self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embed_size, padding_idx=padding_idx)
+        self.lstm = nn.LSTM(input_size=embed_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
+        self.linear = nn.Linear(hidden_size, vocab_size)
+
+        self.init_h = nn.Linear(embed_size, hidden_size)
+        self.init_c = nn.Linear(embed_size, hidden_size)
+    
+    def forward(self, features, captions, hidden=None):
+        if hidden == None:
+            h0 = self.init_h(features).unsqueeze(0).repeat(self.lstm.num_layers, 1, 1)
+            c0 = self.init_c(features).unsqueeze(0).repeat(self.lstm.num_layers, 1, 1)
+            hidden = (h0, c0)
+            # dataflow: (B, seqlen) -> (B, hidden_size) -> (1, B, hidden_size) -> (num_layers, B, hidden_size)
+
+        embeddings = self.embed(captions) # (B, seqlen) ->  Training: (B, seqlen, embed_size) | Inference: (B, 1, embed_size)
+        outputs, hidden = self.lstm(embeddings, hidden) # Training: (B, seqlen, hidden_size) | Inference: (B, 1, hidden_size)
+        outputs = self.linear(outputs) # Training: (B, seqlen, vocab_size) | Inference: (B, 1, vocab_size)
         return outputs, hidden

models/encoder.py → encoder.py

@@ -1,24 +1,24 @@
-from torchvision.models import resnet50, ResNet50_Weights
-import torch.nn as nn
-
-
-class EncoderCNN(nn.Module):
-    def __init__(self, embed_size, fine_tune=False):
-        super(EncoderCNN, self).__init__()
-        resnet = resnet50(weights=ResNet50_Weights.DEFAULT if fine_tune else None)
-        for param in resnet.parameters():
-            param.requires_grad = False
-        if fine_tune:
-            for param in resnet.layer4.parameters():
-                param.requires_grad = True
-        backbone = list(resnet.children())[:-1]
-
-        self.resnet = nn.Sequential(*backbone)
-        self.fc = nn.Linear(resnet.fc.in_features, embed_size)
-        self.bn = nn.BatchNorm1d(num_features=embed_size, momentum=0.01)
-    
-    def forward(self, images): # (B, C, W, H)
-        features = self.resnet(images) # (B, 2048, 1, 1)
-        features = features.reshape(features.shape[0], -1) # (B, 2048*1*1) not necessay to reshape as fc layer can take any size input
-        return self.bn(self.fc(features)) # (B, embed_size)
-
+from torchvision.models import resnet50, ResNet50_Weights
+import torch.nn as nn
+
+
+class EncoderCNN(nn.Module):
+    def __init__(self, embed_size, fine_tune=False):
+        super(EncoderCNN, self).__init__()
+        resnet = resnet50(weights=ResNet50_Weights.DEFAULT if fine_tune else None)
+        for param in resnet.parameters():
+            param.requires_grad = False
+        if fine_tune:
+            for param in resnet.layer4.parameters():
+                param.requires_grad = True
+        backbone = list(resnet.children())[:-1]
+
+        self.resnet = nn.Sequential(*backbone)
+        self.fc = nn.Linear(resnet.fc.in_features, embed_size)
+        self.bn = nn.BatchNorm1d(num_features=embed_size, momentum=0.01)
+    
+    def forward(self, images): # (B, C, W, H)
+        features = self.resnet(images) # (B, 2048, 1, 1)
+        features = features.reshape(features.shape[0], -1) # (B, 2048*1*1) not necessay to reshape as fc layer can take any size input
+        return self.bn(self.fc(features)) # (B, embed_size)
+

utils/helpers.py

@@ -1,15 +1,16 @@
 # from enum import Enum
 import os
+from huggingface_hub import hf_hub_download
 
+repo_id = "Sher1988/image-classifier-weights"
+encoder_path = hf_hub_download(repo_id=repo_id, filename="encoder.pth")
+decoder_path = hf_hub_download(repo_id=repo_id, filename="decoder.pth")
 
-DATA_DIR = 'data'
-CAPTIONS_FILE_PATH = os.path.join(DATA_DIR, 'flickr_data/captions.txt')
-IMAGES_PATH = os.path.join(DATA_DIR, 'flickr_data/Images')
 
+IMAGE_DIR = 'flickr8k/images'
 
-MODELS_PATH = 'models'
-ENCODER_MODEL_PATH = os.path.join(MODELS_PATH, 'encoder.pth')
-DECODER_MODEL_PATH = os.path.join(MODELS_PATH, 'decoder.pth')
-VOCAB_FILE_PATH = os.path.join(MODELS_PATH, 'vocabulary.json')
-TOKANIZED_CAPTIONS = os.path.join(MODELS_PATH, 'captions_tokenized.csv')
-CAPTIONS_TKN_PATH = os.path.join(MODELS_PATH, 'captions_tokenized.csv')
+DATA_DIR = 'data'
+ENCODER_PATH = os.path.join(DATA_DIR, 'encoder.pth')
+DECODER_PATH = os.path.join(DATA_DIR, 'decoder.pth')
+VOCAB_PATH = os.path.join(DATA_DIR, 'vocabulary.json')
+CAPTIONS_PATH = os.path.join(DATA_DIR, 'captions_tokenized.csv')