Update model.py

model.py

@@ -4,13 +4,11 @@ import torchvision.transforms as transforms
 import torchvision.models as models
 from PIL import Image
 import pickle
-import sys
 import os
 import re
 from collections import Counter
 from huggingface_hub import hf_hub_download
 
-
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 EMBED_DIM = 512
@@ -38,6 +36,7 @@ class Vocabulary:
 
     def build_vocabulary(self, sentence_list):
         frequencies = Counter()
+
         for sentence in sentence_list:
             tokens = self.tokenizer(sentence)
             frequencies.update(tokens)
@@ -51,11 +50,10 @@ class Vocabulary:
     def numericalize(self, text):
         tokens = self.tokenizer(text)
         numericalized = []
+
         for token in tokens:
-            if token in self.stoi:
-                numericalized.append(self.stoi[token])
-            else:
-                numericalized.append(self.stoi["unk"])
+            numericalized.append(self.stoi.get(token, self.stoi["unk"]))
+
         return numericalized
 
 
@@ -65,21 +63,28 @@ class Vocabulary:
 class ResNetEncoder(nn.Module):
     def __init__(self, embed_dim):
         super().__init__()
-        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
-        for param in resnet.parameters():
-            param.requires_grad = True
+
+        resnet = models.resnet50(weights=None)
+
         modules = list(resnet.children())[:-1]
+
         self.resnet = nn.Sequential(*modules)
-        
+
         self.fc = nn.Linear(resnet.fc.in_features, embed_dim)
+
         self.batch_norm = nn.BatchNorm1d(embed_dim, momentum=0.01)
 
     def forward(self, images):
+
         with torch.no_grad():
-            features = self.resnet(images)  # (batch_size, 2048, 1, 1)
+            features = self.resnet(images)
+
         features = features.view(features.size(0), -1)
+
         features = self.fc(features)
+
         features = self.batch_norm(features)
+
         return features
 
 
@@ -87,20 +92,31 @@ class ResNetEncoder(nn.Module):
 # Decoder
 # -----------------------
 class DecoderLSTM(nn.Module):
+
     def __init__(self, embed_dim, hidden_dim, vocab_size, num_layers=1):
+
         super().__init__()
+
         self.embedding = nn.Embedding(vocab_size, embed_dim)
+
         self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True)
+
         self.fc = nn.Linear(hidden_dim, vocab_size)
 
     def forward(self, features, captions):
-        # remove the last token for input
-        captions_in = captions[:, :-1]
-        emb = self.embedding(captions_in)
+
+        captions = captions[:, :-1]
+
+        emb = self.embedding(captions)
+
         features = features.unsqueeze(1)
+
         lstm_input = torch.cat((features, emb), dim=1)
+
         outputs, _ = self.lstm(lstm_input)
+
         logits = self.fc(outputs)
+
         return logits
 
 
@@ -108,19 +124,26 @@ class DecoderLSTM(nn.Module):
 # Caption Model
 # -----------------------
 class ImageCaptioningModel(nn.Module):
+
     def __init__(self, encoder, decoder):
+
         super().__init__()
+
         self.encoder = encoder
+
         self.decoder = decoder
 
     def forward(self, images, captions):
+
         features = self.encoder(images)
+
         outputs = self.decoder(features, captions)
+
         return outputs
 
 
 # -----------------------
-# Caption generator
+# Caption Generator
 # -----------------------
 def generate_caption(model, image, vocab):
 
@@ -128,102 +151,103 @@ def generate_caption(model, image, vocab):
 
     image = image.unsqueeze(0).to(DEVICE)
 
+    sentence = []
+
     with torch.no_grad():
-        # Get image features
-        features = model.encoder(image)  # (1, embed_dim)
 
-        # Start with the start token
+        features = model.encoder(image)
+
         word_idx = vocab.stoi["startofseq"]
-        sentence = []
-        
-        # Initialize hidden state for LSTM
-        h = None
-        
+
+        hidden = None
+
         for _ in range(MAX_LEN):
-            # Create input: concatenate features with embedding of previous word
+
             word_tensor = torch.tensor([word_idx]).to(DEVICE)
-            emb = model.decoder.embedding(word_tensor)  # (1, embed_dim)
-            
-            if h is None:
-                # First step: concatenate features with embedding
-                lstm_input = torch.cat([features.unsqueeze(1), emb.unsqueeze(1)], dim=1)  # (1, 2, embed_dim)
+
+            emb = model.decoder.embedding(word_tensor)
+
+            if hidden is None:
+
+                lstm_input = torch.cat(
+                    [features.unsqueeze(1), emb.unsqueeze(1)], dim=1
+                )
+
             else:
-                lstm_input = emb.unsqueeze(1)  # (1, 1, embed_dim)
-            
-            # Forward through LSTM
-            output, h_new = model.decoder.lstm(lstm_input, h)
-            h = h_new
-            
-            # Predict next token
-            logits = model.decoder.fc(output[:, -1, :])  # (1, vocab_size)
+
+                lstm_input = emb.unsqueeze(1)
+
+            output, hidden = model.decoder.lstm(lstm_input, hidden)
+
+            logits = model.decoder.fc(output[:, -1, :])
+
             predicted = logits.argmax(1).item()
-            
-            # Get token from vocab
+
             token = vocab.itos[predicted]
-            
+
             if token == "endofseq":
                 break
-            
+
             sentence.append(token)
+
             word_idx = predicted
 
     return " ".join(sentence)
 
 
 # -----------------------
-# Image transform
+# Image Transform
 # -----------------------
-transform = transforms.Compose([
-    transforms.Resize((224,224)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485,0.456,0.406],
-        std=[0.229,0.224,0.225]
-    )
-])
+transform = transforms.Compose(
+    [
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ),
+    ]
+)
 
 
 # -----------------------
-# Main
+# Load Model Once
 # -----------------------
-def main():
 
-    image_path = sys.argv[1]
+script_dir = os.path.dirname(os.path.abspath(__file__))
 
-    # Get the directory where this script is located
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    CHECKPOINT_PATH = hf_hub_download(
+CHECKPOINT_PATH = hf_hub_download(
     repo_id="VIKRAM989/image-label",
     filename="best_checkpoint.pth"
-    )
-    VOCAB_PATH = os.path.join(script_dir, "vocab.pkl")
+)
 
-    # load vocab
-    with open(VOCAB_PATH, "rb") as f:
-        vocab = pickle.load(f)
+VOCAB_PATH = os.path.join(script_dir, "vocab.pkl")
 
-    vocab_size = len(vocab)
+with open(VOCAB_PATH, "rb") as f:
+    vocab = pickle.load(f)
 
-    # rebuild model
-    encoder = ResNetEncoder(EMBED_DIM)
-    decoder = DecoderLSTM(EMBED_DIM, HIDDEN_DIM, vocab_size)
-    model = ImageCaptioningModel(encoder, decoder).to(DEVICE)
+vocab_size = len(vocab)
 
-    # load checkpoint
-    checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+encoder = ResNetEncoder(EMBED_DIM)
 
-    model.load_state_dict(checkpoint["model_state_dict"])
+decoder = DecoderLSTM(EMBED_DIM, HIDDEN_DIM, vocab_size)
 
-    model.eval()
+model = ImageCaptioningModel(encoder, decoder).to(DEVICE)
 
-    # load image
-    img = Image.open(image_path).convert("RGB")
-    img = transform(img)
+checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
 
-    caption = generate_caption(model, img, vocab)
+model.load_state_dict(checkpoint["model_state_dict"])
+
+model.eval()
 
-    print("\nCaption:", caption)
 
+# -----------------------
+# Public Function for API
+# -----------------------
+def caption_image(pil_image):
+
+    img = transform(pil_image).to(DEVICE)
+
+    caption = generate_caption(model, img, vocab)
 
-if __name__ == "__main__":
-    main()
+    return caption