Create app.py

app.py

@@ -0,0 +1,165 @@
+import torch
+import torch.nn as nn
+import gradio as gr
+from PIL import Image
+import torchvision.transforms as transforms
+from torchvision import models
+
+# -------------------------
+# Model definitions (must match training)
+# -------------------------
+PAD_TOKEN = "<pad>"
+UNK_TOKEN = "<unk>"
+
+class Encoder(nn.Module):
+    def __init__(self, in_dim=2048, hidden_size=512):
+        super().__init__()
+        self.fc = nn.Linear(in_dim, hidden_size)
+        self.relu = nn.ReLU()
+
+    def forward(self, feat):
+        return self.relu(self.fc(feat))
+
+class Decoder(nn.Module):
+    def __init__(self, vocab_size, pad_id, embed_dim=256, hidden_size=512, dropout=0.1):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_id)
+        self.lstm = nn.LSTM(embed_dim, hidden_size, num_layers=1, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+        self.fc_out = nn.Linear(hidden_size, vocab_size)
+
+class Img2Caption(nn.Module):
+    def __init__(self, vocab_size, pad_id, hidden_size=512, embed_dim=256):
+        super().__init__()
+        self.encoder = Encoder(in_dim=2048, hidden_size=hidden_size)
+        self.decoder = Decoder(vocab_size=vocab_size, pad_id=pad_id, embed_dim=embed_dim, hidden_size=hidden_size)
+
+# -------------------------
+# Load checkpoint
+# -------------------------
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CKPT_PATH = "img_caption_seq2seq.pth"
+
+ckpt = torch.load(CKPT_PATH, map_location=DEVICE)
+word2idx = ckpt["word2idx"]
+idx2word = ckpt["idx2word"]
+max_len = ckpt.get("max_len", 30)
+
+pad_id = word2idx[PAD_TOKEN]
+start_id = word2idx["<start>"]
+end_id = word2idx["<end>"]
+
+model = Img2Caption(vocab_size=len(word2idx), pad_id=pad_id).to(DEVICE)
+model.load_state_dict(ckpt["model_state"])
+model.eval()
+
+# -------------------------
+# ResNet50 feature extractor (on-the-fly)
+# -------------------------
+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 decode_tokens(token_ids):
+    words = []
+    for tid in token_ids:
+        w = idx2word.get(int(tid), UNK_TOKEN)
+        if w == "<end>":
+            break
+        if w not in ["<start>", "<pad>"]:
+            words.append(w)
+    return " ".join(words)
+
+@torch.no_grad()
+def greedy_caption(feat_vec, max_words=30):
+    feat = torch.tensor(feat_vec, dtype=torch.float32).unsqueeze(0).to(DEVICE)  # [1,2048]
+    h0 = model.encoder(feat)  # [1,hidden]
+
+    last = start_id
+    out_tokens = []
+
+    h = h0.unsqueeze(0)  # [1,1,hidden]
+    c = torch.zeros_like(h)
+
+    for _ in range(max_words):
+        cur = torch.tensor([[last]], dtype=torch.long).to(DEVICE)
+        emb = model.decoder.embed(cur)  # [1,1,E]
+        lstm_out, (h, c) = model.decoder.lstm(emb, (h, c))  # [1,1,H]
+        logits = model.decoder.fc_out(lstm_out.squeeze(1))  # [1,V]
+        nxt = int(torch.argmax(logits, dim=-1).item())
+
+        if nxt == end_id:
+            break
+        out_tokens.append(nxt)
+        last = nxt
+
+    return decode_tokens(out_tokens)
+
+@torch.no_grad()
+def beam_caption(feat_vec, beam_size=3, max_words=30):
+    feat = torch.tensor(feat_vec, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+    h0 = model.encoder(feat)
+
+    h = h0.unsqueeze(0)
+    c = torch.zeros_like(h)
+
+    beams = [([], 0.0, h, c, start_id)]  # (tokens, score, h, c, last)
+
+    for _ in range(max_words):
+        new_beams = []
+        for tokens, score, h_i, c_i, last in beams:
+            if last == end_id:
+                new_beams.append((tokens, score, h_i, c_i, last))
+                continue
+
+            cur = torch.tensor([[last]], dtype=torch.long).to(DEVICE)
+            emb = model.decoder.embed(cur)
+            lstm_out, (h_new, c_new) = model.decoder.lstm(emb, (h_i, c_i))
+            logits = model.decoder.fc_out(lstm_out.squeeze(1))
+            log_probs = torch.log_softmax(logits, dim=-1).squeeze(0)
+
+            topk = torch.topk(log_probs, beam_size)
+            for lp, idx in zip(topk.values.tolist(), topk.indices.tolist()):
+                new_beams.append((tokens + [idx], score + lp, h_new, c_new, idx))
+
+        new_beams.sort(key=lambda x: x[1], reverse=True)
+        beams = new_beams[:beam_size]
+
+        if all(b[4] == end_id for b in beams):
+            break
+
+    best = beams[0][0]
+    if len(best) and best[-1] == end_id:
+        best = best[:-1]
+    return decode_tokens(best)
+
+@torch.no_grad()
+def caption_image(img: Image.Image, decoding="Beam Search"):
+    img = img.convert("RGB")
+    x = transform(img).unsqueeze(0).to(DEVICE)
+
+    feat = resnet(x).view(1, -1).squeeze(0).cpu().numpy()  # [2048]
+
+    if decoding == "Greedy":
+        return greedy_caption(feat, max_words=30)
+    return beam_caption(feat, beam_size=3, max_words=30)
+
+demo = gr.Interface(
+    fn=caption_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(["Beam Search", "Greedy"], value="Beam Search", label="Decoding")
+    ],
+    outputs=gr.Textbox(label="Generated Caption"),
+    title="Seq2Seq Image Captioning (Flickr30k)",
+    description="Upload an image and generate a caption using a ResNet50 + LSTM Seq2Seq model."
+)
+
+if __name__ == "__main__":
+    demo.launch()