Spaces:
Build error
Build error
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import pickle | |
| import numpy as np | |
| from PIL import Image | |
| from torchvision import models, transforms | |
| # ============================================================================== | |
| # EXACT MODEL ARCHITECTURE — copied 1-to-1 from training code | |
| # ============================================================================== | |
| class Encoder(nn.Module): | |
| def __init__(self, feature_size=2048, hidden_size=512, dropout=0.5): | |
| super(Encoder, self).__init__() | |
| self.fc = nn.Linear(feature_size, hidden_size) | |
| self.dropout = nn.Dropout(dropout) | |
| self.bn = nn.BatchNorm1d(hidden_size) | |
| def forward(self, features): | |
| encoded = self.fc(features) | |
| encoded = self.bn(encoded) | |
| encoded = F.relu(encoded) | |
| encoded = self.dropout(encoded) | |
| return encoded | |
| class Attention(nn.Module): | |
| def __init__(self, encoder_dim, decoder_dim, attention_dim): | |
| super(Attention, self).__init__() | |
| self.encoder_att = nn.Linear(encoder_dim, attention_dim) | |
| self.decoder_att = nn.Linear(decoder_dim, attention_dim) | |
| self.full_att = nn.Linear(attention_dim, 1) | |
| self.relu = nn.ReLU() | |
| self.softmax = nn.Softmax(dim=1) | |
| def forward(self, encoder_out, decoder_hidden): | |
| att1 = self.encoder_att(encoder_out) | |
| att2 = self.decoder_att(decoder_hidden) | |
| att = self.full_att(self.relu(att1 + att2)).squeeze(-1) | |
| alpha = self.softmax(att.unsqueeze(1)) | |
| attention_weighted_encoding = encoder_out * alpha | |
| return attention_weighted_encoding, alpha | |
| class Decoder(nn.Module): | |
| def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1, dropout=0.5): | |
| super(Decoder, self).__init__() | |
| self.embed_size = embed_size | |
| self.hidden_size = hidden_size | |
| self.vocab_size = vocab_size | |
| self.num_layers = num_layers | |
| self.embedding = nn.Embedding(vocab_size, embed_size) | |
| self.dropout = nn.Dropout(dropout) | |
| self.lstm = nn.LSTM( | |
| embed_size + hidden_size, | |
| hidden_size, | |
| num_layers, | |
| batch_first=True, | |
| dropout=dropout if num_layers > 1 else 0 | |
| ) | |
| self.attention = Attention(hidden_size, hidden_size, hidden_size) | |
| self.fc = nn.Linear(hidden_size, vocab_size) | |
| self._init_weights() | |
| def _init_weights(self): | |
| self.embedding.weight.data.uniform_(-0.1, 0.1) | |
| self.fc.bias.data.fill_(0) | |
| self.fc.weight.data.uniform_(-0.1, 0.1) | |
| def generate(self, features, max_len=50, start_token=1, end_token=2): | |
| """Greedy search""" | |
| h = features.unsqueeze(0).repeat(self.num_layers, 1, 1) | |
| c = torch.zeros_like(h) | |
| input_token = torch.LongTensor([start_token]).to(features.device) | |
| generated = [] | |
| for _ in range(max_len): | |
| embed = self.embedding(input_token).unsqueeze(1) | |
| hiddens = h[-1] | |
| context, _ = self.attention(features, hiddens) | |
| lstm_input = torch.cat([embed.squeeze(1), context], dim=1).unsqueeze(1) | |
| _, (h, c) = self.lstm(lstm_input, (h, c)) | |
| output = self.fc(h[-1]) | |
| predicted = output.argmax(1) | |
| generated.append(predicted.item()) | |
| if predicted.item() == end_token: | |
| break | |
| input_token = predicted | |
| return generated | |
| def beam_search(self, features, beam_width=5, max_len=50, start_token=1, end_token=2): | |
| """Beam search""" | |
| device = features.device | |
| h = features.unsqueeze(0).repeat(self.num_layers, 1, 1) | |
| c = torch.zeros_like(h) | |
| beams = [([start_token], 0.0, h, c)] | |
| completed_beams = [] | |
| for _ in range(max_len): | |
| all_candidates = [] | |
| for seq, score, h_state, c_state in beams: | |
| if seq[-1] == end_token: | |
| completed_beams.append((seq, score)) | |
| continue | |
| input_token = torch.LongTensor([seq[-1]]).to(device) | |
| embed = self.embedding(input_token).unsqueeze(1) | |
| hiddens = h_state[-1] | |
| context, _ = self.attention(features, hiddens) | |
| lstm_input = torch.cat([embed.squeeze(1), context], dim=1).unsqueeze(1) | |
| _, (h_new, c_new) = self.lstm(lstm_input, (h_state, c_state)) | |
| log_probs = F.log_softmax(self.fc(h_new[-1]), dim=1) | |
| topk_probs, topk_idx = log_probs.topk(beam_width, dim=1) | |
| for i in range(beam_width): | |
| all_candidates.append(( | |
| seq + [topk_idx[0, i].item()], | |
| score + topk_probs[0, i].item(), | |
| h_new, c_new | |
| )) | |
| beams = sorted(all_candidates, key=lambda x: x[1], reverse=True)[:beam_width] | |
| if not beams: | |
| break | |
| completed_beams.extend([(seq, score) for seq, score, _, _ in beams]) | |
| if completed_beams: | |
| best_seq, _ = max(completed_beams, key=lambda x: x[1]) | |
| return best_seq | |
| return beams[0][0] if beams else [start_token, end_token] | |
| class ImageCaptioningModel(nn.Module): | |
| def __init__(self, encoder, decoder): | |
| super(ImageCaptioningModel, self).__init__() | |
| self.encoder = encoder | |
| self.decoder = decoder | |
| def forward(self, features, captions, lengths): | |
| return self.decoder(self.encoder(features), captions, lengths) | |
| # ============================================================================== | |
| # LOAD MODELS & VOCAB (runs once at startup) | |
| # ============================================================================== | |
| DEVICE = torch.device("cpu") # HF free-tier has no GPU | |
| # ── Hyperparameters — must match training exactly ────────────────────────────── | |
| FEATURE_SIZE = 2048 | |
| HIDDEN_SIZE = 512 | |
| EMBED_SIZE = 512 | |
| NUM_LAYERS = 2 | |
| DROPOUT = 0.5 | |
| def load_vocab(path="vocab.pkl"): | |
| with open(path, "rb") as f: | |
| return pickle.load(f) | |
| def load_model(vocab, path="best_model.pth"): | |
| encoder = Encoder(feature_size=FEATURE_SIZE, hidden_size=HIDDEN_SIZE, dropout=DROPOUT) | |
| decoder = Decoder( | |
| embed_size=EMBED_SIZE, | |
| hidden_size=HIDDEN_SIZE, | |
| vocab_size=len(vocab), | |
| num_layers=NUM_LAYERS, | |
| dropout=DROPOUT | |
| ) | |
| model = ImageCaptioningModel(encoder, decoder).to(DEVICE) | |
| checkpoint = torch.load(path, map_location=DEVICE, weights_only=False) | |
| model.load_state_dict(checkpoint["model_state_dict"]) | |
| model.eval() | |
| return model | |
| def load_resnet(): | |
| resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) | |
| resnet = nn.Sequential(*list(resnet.children())[:-1]) # drop classifier | |
| resnet = resnet.to(DEVICE) | |
| resnet.eval() | |
| return resnet | |
| print("Loading vocab...") | |
| vocab = load_vocab("vocab.pkl") | |
| print(f" Vocab size: {len(vocab)} words") | |
| print("Loading captioning model...") | |
| model = load_model(vocab, "best_model.pth") | |
| print(" Model ready.") | |
| print("Loading ResNet50 feature extractor...") | |
| resnet = load_resnet() | |
| print(" ResNet50 ready.") | |
| # ── Image pre-processing — same as training ──────────────────────────────────── | |
| IMG_TRANSFORM = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), | |
| ]) | |
| # ============================================================================== | |
| # INFERENCE HELPERS | |
| # ============================================================================== | |
| def extract_features(pil_image: Image.Image) -> torch.Tensor: | |
| """Run image through ResNet50 → (1, 2048) feature vector.""" | |
| img_tensor = IMG_TRANSFORM(pil_image).unsqueeze(0).to(DEVICE) | |
| with torch.no_grad(): | |
| features = resnet(img_tensor).view(1, -1) # (1, 2048) | |
| return features | |
| def ids_to_caption(ids, vocab) -> str: | |
| """Convert list of token ids → clean sentence string.""" | |
| words = [] | |
| for idx in ids: | |
| if idx == vocab.stoi["<end>"]: | |
| break | |
| if idx not in (vocab.stoi["<start>"], vocab.stoi["<pad>"]): | |
| words.append(vocab.itos[idx]) | |
| return " ".join(words).capitalize() if words else "Could not generate a caption." | |
| def caption_image(pil_image: Image.Image, method: str, beam_width: int) -> tuple[str, str]: | |
| """ | |
| Main inference function called by Gradio. | |
| Returns (caption_string, info_string). | |
| """ | |
| if pil_image is None: | |
| return "Please upload an image.", "" | |
| # 1. Extract features | |
| features = extract_features(pil_image) | |
| # 2. Encode | |
| with torch.no_grad(): | |
| encoded = model.encoder(features) # (1, 512) | |
| # 3. Decode | |
| with torch.no_grad(): | |
| if method == "Beam Search": | |
| ids = model.decoder.beam_search( | |
| encoded, | |
| beam_width=beam_width, | |
| max_len=50, | |
| start_token=vocab.stoi["<start>"], | |
| end_token=vocab.stoi["<end>"], | |
| ) | |
| else: # Greedy | |
| ids = model.decoder.generate( | |
| encoded, | |
| max_len=50, | |
| start_token=vocab.stoi["<start>"], | |
| end_token=vocab.stoi["<end>"], | |
| ) | |
| caption = ids_to_caption(ids, vocab) | |
| info = f"Method: {method}" + (f" | Beam width: {beam_width}" if method == "Beam Search" else "") | |
| return caption, info | |
| # ============================================================================== | |
| # GRADIO UI | |
| # ============================================================================== | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Playfair+Display:ital,wght@0,700;1,400&family=DM+Mono:wght@400;500&display=swap'); | |
| :root { | |
| --cream: #F5F0E8; | |
| --ink: #1A1208; | |
| --sienna: #B85C2A; | |
| --gold: #D4A843; | |
| --warm-gray: #8C7B6B; | |
| --border: #D9CFC2; | |
| } | |
| body, .gradio-container { | |
| background: var(--cream) !important; | |
| font-family: 'DM Mono', monospace !important; | |
| } | |
| /* ── Header ── */ | |
| #header { | |
| text-align: center; | |
| padding: 2.5rem 1rem 1rem; | |
| border-bottom: 2px solid var(--ink); | |
| margin-bottom: 2rem; | |
| } | |
| #header h1 { | |
| font-family: 'Playfair Display', serif !important; | |
| font-size: clamp(2rem, 5vw, 3.5rem) !important; | |
| font-weight: 700 !important; | |
| color: var(--ink) !important; | |
| letter-spacing: -0.02em; | |
| margin: 0 !important; | |
| line-height: 1.1; | |
| } | |
| #header .subtitle { | |
| font-family: 'Playfair Display', serif; | |
| font-style: italic; | |
| font-size: 1.1rem; | |
| color: var(--sienna); | |
| margin-top: 0.4rem; | |
| } | |
| #header .tag-line { | |
| font-size: 0.72rem; | |
| letter-spacing: 0.18em; | |
| text-transform: uppercase; | |
| color: var(--warm-gray); | |
| margin-top: 0.6rem; | |
| } | |
| /* ── Main columns ── */ | |
| .main-col { padding: 0 0.5rem; } | |
| /* ── Image upload box ── */ | |
| .image-upload .wrap { border: 2px dashed var(--border) !important; border-radius: 4px !important; background: white !important; } | |
| .image-upload .wrap:hover { border-color: var(--sienna) !important; } | |
| /* ── Caption output ── */ | |
| #caption-box textarea { | |
| font-family: 'Playfair Display', serif !important; | |
| font-style: italic !important; | |
| font-size: 1.35rem !important; | |
| color: var(--ink) !important; | |
| background: white !important; | |
| border: 2px solid var(--ink) !important; | |
| border-radius: 4px !important; | |
| padding: 1.2rem !important; | |
| line-height: 1.6 !important; | |
| min-height: 80px !important; | |
| } | |
| /* ── Info box ── */ | |
| #info-box textarea { | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 0.78rem !important; | |
| color: var(--warm-gray) !important; | |
| background: transparent !important; | |
| border: none !important; | |
| border-top: 1px solid var(--border) !important; | |
| padding: 0.6rem 0 0 !important; | |
| } | |
| /* ── Controls ── */ | |
| .gr-radio label, .gr-slider label { | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 0.8rem !important; | |
| letter-spacing: 0.05em !important; | |
| color: var(--ink) !important; | |
| } | |
| /* ── Generate button ── */ | |
| #generate-btn { | |
| background: var(--ink) !important; | |
| color: var(--cream) !important; | |
| border: none !important; | |
| border-radius: 2px !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 0.85rem !important; | |
| letter-spacing: 0.12em !important; | |
| text-transform: uppercase !important; | |
| padding: 0.75rem 2rem !important; | |
| cursor: pointer !important; | |
| transition: background 0.2s, transform 0.1s !important; | |
| width: 100% !important; | |
| } | |
| #generate-btn:hover { | |
| background: var(--sienna) !important; | |
| transform: translateY(-1px) !important; | |
| } | |
| /* ── Section labels ── */ | |
| .section-label { | |
| font-size: 0.68rem; | |
| letter-spacing: 0.18em; | |
| text-transform: uppercase; | |
| color: var(--warm-gray); | |
| border-bottom: 1px solid var(--border); | |
| padding-bottom: 0.4rem; | |
| margin-bottom: 0.8rem; | |
| } | |
| /* ── Examples strip ── */ | |
| .examples-label { | |
| font-size: 0.68rem; | |
| letter-spacing: 0.18em; | |
| text-transform: uppercase; | |
| color: var(--warm-gray); | |
| margin-top: 1.5rem; | |
| margin-bottom: 0.5rem; | |
| } | |
| /* ── Footer ── */ | |
| #footer { | |
| text-align: center; | |
| padding: 1.5rem; | |
| border-top: 1px solid var(--border); | |
| margin-top: 2rem; | |
| font-size: 0.72rem; | |
| letter-spacing: 0.1em; | |
| color: var(--warm-gray); | |
| } | |
| /* hide default Gradio labels where we use custom ones */ | |
| .hide-label > label { display: none !important; } | |
| """ | |
| def run(image, method, beam_width): | |
| return caption_image(image, method, int(beam_width)) | |
| with gr.Blocks(css=CSS, title="Neural Storyteller") as demo: | |
| # ── Header ──────────────────────────────────────────────────────────────── | |
| gr.HTML(""" | |
| <div id="header"> | |
| <div class="tag-line">ResNet50 · LSTM · Bahdanau Attention · Flickr30k</div> | |
| <h1>Neural Storyteller</h1> | |
| <div class="subtitle">Image Captioning with Seq2Seq</div> | |
| </div> | |
| """) | |
| # ── Main layout ─────────────────────────────────────────────────────────── | |
| with gr.Row(): | |
| # Left — image upload | |
| with gr.Column(scale=5, elem_classes="main-col"): | |
| gr.HTML('<div class="section-label">Upload Image</div>') | |
| image_input = gr.Image( | |
| type="pil", | |
| label="", | |
| elem_classes=["image-upload", "hide-label"], | |
| height=380, | |
| ) | |
| # Right — controls + output | |
| with gr.Column(scale=5, elem_classes="main-col"): | |
| gr.HTML('<div class="section-label">Decoding Strategy</div>') | |
| method = gr.Radio( | |
| choices=["Beam Search", "Greedy"], | |
| value="Beam Search", | |
| label="", | |
| elem_classes="hide-label", | |
| ) | |
| beam_width = gr.Slider( | |
| minimum=1, maximum=10, value=5, step=1, | |
| label="Beam Width (only applies to Beam Search)", | |
| ) | |
| gr.HTML('<div class="section-label" style="margin-top:1.2rem">Generated Caption</div>') | |
| caption_out = gr.Textbox( | |
| label="", | |
| elem_id="caption-box", | |
| elem_classes="hide-label", | |
| lines=3, | |
| interactive=False, | |
| placeholder="Your caption will appear here…", | |
| ) | |
| info_out = gr.Textbox( | |
| label="", | |
| elem_id="info-box", | |
| elem_classes="hide-label", | |
| lines=1, | |
| interactive=False, | |
| ) | |
| generate_btn = gr.Button("Generate Caption", elem_id="generate-btn") | |
| # ── Wire up ─────────────────────────────────────────────────────────────── | |
| generate_btn.click( | |
| fn=run, | |
| inputs=[image_input, method, beam_width], | |
| outputs=[caption_out, info_out], | |
| ) | |
| # Also trigger on image upload for instant feel | |
| image_input.upload( | |
| fn=run, | |
| inputs=[image_input, method, beam_width], | |
| outputs=[caption_out, info_out], | |
| ) | |
| # ── Footer ──────────────────────────────────────────────────────────────── | |
| gr.HTML(""" | |
| <div id="footer"> | |
| Neural Storyteller · Seq2Seq with Bahdanau Attention | |
| · Trained on Flickr30k · Built with PyTorch | |
| </div> | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch() | |