Building a Neural Storyteller: Image Captioning with Seq2Seq Models

Introduction

In the realm of Generative AI, combining computer vision and natural language processing opens up exciting possibilities. As part of our Assignment №1 for the course AI4009 (Generative AI) at the National University of Computer and Emerging Sciences (NUCES) in Spring 2026, we Jaweria Shahid and Aliaworked as a group to build a multimodal deep learning model called Neural Storyteller. This project uses a Sequence-to-Sequence (Seq2Seq) architecture to generate natural language descriptions for images.

We leveraged the Flickr30k dataset on Kaggle, employing a pre-trained ResNet50 for image feature extraction and a custom Seq2Seq model with an LSTM decoder for caption generation. The goal was to create an end-to-end system that processes images, trains on captions, evaluates performance, and deploys via a Gradio app. In this post, we’ll highlight 2–3 major parts of our code implementation, focusing on the core technical aspects without diving into every detail.

Part 1: Feature Extraction Pipeline

Training a CNN from scratch alongside an RNN can be computationally intensive, so we “cached” image features upfront using a pre-trained ResNet50 model. This converts each image into a 2048-dimensional feature vector, which serves as input to our Seq2Seq encoder.

We set up the environment on Kaggle with GPU acceleration (T4 x2) and loaded the Flickr30k dataset. The code walks through the input directory to locate images, applies transformations (resize, normalize), and extracts features in batches for efficiency. The results are saved to a pickle file (flickr30k_features.pkl) for quick reuse.

Here’s the key code snippet for this pipeline:

Python

import os, pickle, torch, torch.nn as nn
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
from tqdm import tqdm

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model = nn.Sequential(*list(model.children())[:-1])  # Feature vector only
model = nn.DataParallel(model).to(device)
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

dataset = FlickrDataset(IMAGE_DIR, transform)
loader = DataLoader(dataset, batch_size=128, num_workers=4)
features_dict = {}

with torch.no_grad():
    for imgs, names in tqdm(loader, desc="Extracting Features"):
        feats = model(imgs.to(device)).view(imgs.size(0), -1)
        for i, name in enumerate(names):
            features_dict[name] = feats[i].cpu().numpy()

with open(OUTPUT_FILE, 'wb') as f:
    pickle.dump(features_dict, f)

print(f"Success! {len(features_dict)} images processed and saved to {OUTPUT_FILE}")

This step processed over 31,000 images efficiently, caching features to avoid redundant computations during training.

Part 2: The Seq2Seq Architecture and Training

The heart of the model is the Seq2Seq setup: an Encoder that projects the image features into a hidden state, and a Decoder (LSTM) that generates captions word-by-word. We preprocessed captions by cleaning text, tokenizing with <start> and <end> tokens, and building a vocabulary (size ~7,691 words, filtered for frequency >=5).

For training, we used CrossEntropyLoss (ignoring padding), Adam optimizer, and implemented both Greedy and Beam Search for inference. The model was trained over epochs, with validation to monitor loss and save the best checkpoint.

Key code for the model architecture and training loop:

Python

class CaptionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

def forward(self, feats, caps, lens):
        hidden = self.encoder(feats).unsqueeze(0)
        hidden = (hidden, torch.zeros_like(hidden))  # For LSTM cell state
        outputs, _ = self.decoder(caps, hidden)
        return outputs

# Training setup (simplified)
model = CaptionModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss(ignore_index=0)  # Ignore padding

# Example training loop
for epoch in range(10):  # EPOCHS
    model.train()
    total_loss = 0
    for feats, caps, lens in train_loader:
        feats, caps = feats.to(device), caps.to(device)
        outputs = model(feats, caps[:, :-1], lens-1)
        loss = criterion(outputs.view(-1, vocab_size), caps[:, 1:].reshape(-1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}: Avg Loss = {total_loss / len(train_loader):.4f}")

# Save best model
torch.save(model.state_dict(), 'best_model.pth')

We also implemented Beam Search for better caption quality during inference, achieving a BLEU-4 score around 0.15–0.20 on validation.

Part 3: Evaluation and Deployment

To evaluate, we plotted training/validation loss curves, computed BLEU-4, Precision/Recall/F1 (token-level), and optionally METEOR/ROUGE. For deployment, we built a Gradio app that takes an image upload, generates a caption, and displays ground truth if it’s from the dataset.

Here’s a snippet for the Gradio interface:

Python

def predict(img):
    # Process image, extract features, generate caption using beam search
    generated_caption = generate_beam(extract_features(img))  # Custom functions
    return f"Model Generated: {generated_caption}"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Textbox(label="Caption Output"),
    title="Neural Storyteller – Image Captioning"
)
demo.launch()

This app makes the model interactive, showcasing its real-world potential.

Conclusion

This project taught us the intricacies of multimodal AI, from feature caching to Seq2Seq training. While challenges like overfitting and caption diversity arose, the results were promising. As Alia and Javeria, we’re excited to explore more in Generative AI. Check our GitHub repo for the full code: https://github.com/Alia-raza/Gen-Ai . Feedback welcome!

Note: This post is based on our group assignment submission.