Upload 4 files

app.py

@@ -0,0 +1,126 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import pickle
+from torchvision import models, transforms
+from PIL import Image
+
+class Config:
+    embed_size = 300
+    hidden_size = 512
+    num_layers = 1
+    feature_dim = 2048
+
+class Encoder(nn.Module):
+    def __init__(self, input_dim, hidden_dim):
+        super(Encoder, self).__init__()
+        self.linear = nn.Linear(input_dim, hidden_dim)
+        self.bn = nn.BatchNorm1d(hidden_dim)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, images):
+        x = self.linear(images)
+        x = self.bn(x)
+        return self.dropout(self.relu(x))
+
+class Decoder(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers):
+        super(Decoder, self).__init__()
+        self.embed = nn.Embedding(vocab_size, embed_size)
+        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
+        self.linear = nn.Linear(hidden_size, vocab_size)
+
+    def forward(self, features, captions):
+        return None
+
+class Seq2Seq(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, feature_dim):
+        super(Seq2Seq, self).__init__()
+        self.encoder = Encoder(feature_dim, hidden_size)
+        self.decoder = Decoder(embed_size, hidden_size, vocab_size, num_layers)
+
+device = torch.device("cpu")
+
+with open('vocab_safe.pkl', 'rb') as f:
+    vocab_data = pickle.load(f)
+itos = vocab_data['itos']
+stoi = vocab_data['stoi']
+vocab_size = len(itos)
+
+model = Seq2Seq(Config.embed_size, Config.hidden_size, vocab_size, Config.num_layers, Config.feature_dim)
+model.load_state_dict(torch.load('best_model.pth', map_location=device))
+model.eval()
+
+resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+resnet = nn.Sequential(*list(resnet.children())[:-1]).to(device)
+resnet.eval()
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+])
+
+def generate_caption(image):
+    try:
+        if image is None:
+            return "Please upload an image first."
+
+        image = image.convert('RGB')
+        img_tensor = transform(image).unsqueeze(0).to(device)
+        
+        with torch.no_grad():
+            features = resnet(img_tensor).view(1, -1) 
+        
+        with torch.no_grad():
+            enc_out = model.encoder(features).unsqueeze(0) 
+            h, c = enc_out, enc_out
+            
+            word_idx = stoi['<start>']
+            word = torch.tensor(word_idx).view(1).to(device)
+            caption = []
+            
+            for i in range(20):
+                embed = model.decoder.embed(word).view(1, 1, -1)
+                output, (h, c) = model.decoder.lstm(embed, (h, c))
+                prediction = model.decoder.linear(output)
+                idx = prediction.argmax(2).item()
+                
+                if idx == stoi['<end>']:
+                    break
+                
+                word_str = itos.get(idx, "<unk>")
+                caption.append(word_str)
+                word = torch.tensor(idx).view(1).to(device)
+                
+        final_caption = " ".join(caption).strip().capitalize()
+        if final_caption:
+            final_caption += "."
+        return final_caption
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🖼️ Image Captioning Generator
+        Upload an image to generate a descriptive caption.
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Image")
+            generate_btn = gr.Button("✨ Generate Caption", variant="primary")
+        with gr.Column():
+            caption_output = gr.Textbox(label="Generated Caption", lines=4, interactive=False)
+            
+    generate_btn.click(
+        fn=generate_caption,
+        inputs=image_input,
+        outputs=caption_output
+    )
+
+if __name__ == "__main__":
+    demo.launch()

best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88cf3415474790fc2aadfc6906bddd4a91a85ddccf536d70fb09cc6b8c40e01c
+size 51560373

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+pillow
+gradio

vocab_safe.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3787a6494232df5dcfa088f6b8d5efbd9e4f23507c0079acb26e85989e67c967
+size 260287