An end-to-end deep learning system that generates natural language descriptions for images. Combines a pre-trained VGG16 convolutional network for visual feature extraction with an LSTM-based language decoder — built with PyTorch.
- Encoder-Decoder Architecture — VGG16 image encoder + LSTM text decoder fused via feature addition
- BLEU Evaluation — Standard MT metrics (BLEU-1, BLEU-2) for caption quality assessment
- Flickr8k Dataset — 8,000 images with 5 human-annotated captions each
- Kaggle-Ready — Designed for GPU training on Kaggle (Tesla T4 compatible)
- Modular Data Pipeline — On-the-fly data generator with batch processing for memory efficiency
Model architecture: VGG16 feature extractor → Dense projection → LSTM decoder → Softmax output
Input Image → [VGG16] → Image Features (4096)
↓
Caption Seed → [Embedding → LSTM] → Feature Fusion → Dense → Softmax → Next Word
Sample Output — Given an unseen image, the model generates:
| Image | Generated Caption | Ground Truth |
|---|---|---|
| (example) | "a dog is running through the grass" | "a brown dog runs across a grassy field" |
| (example) | "a group of children playing in the water" | "kids are playing in a pool of water" |
- Pre-trained on ImageNet (14M+ images) via
torchvision.models - Truncated at the second fully-connected layer (fc2 + ReLU)
- Output: 4,096-dimensional image feature vector
- No fine-tuning — features extracted once, cached to disk with
torch.save
- Embedding Layer: Maps tokenized words to 256-dimensional dense vectors
- LSTM Layer: 256 units — learns sequential dependencies in caption text
- Feature Fusion: Image features (through Dense(256) + LeakyReLU) added to LSTM output
- Output Layer: Dense + Softmax over the vocabulary (vocabulary size)
Image Features (4096)
↓
Dropout(0.3) → Linear(256) → LeakyReLU
↓
Caption Tokens → Embedding(256) → Dropout(0.3) → LSTM(256)
↓
Add ←———————
↓
Linear(256) → LeakyReLU
↓
Linear(vocab_size) → Softmax
| Component | Choice | Rationale |
|---|---|---|
| Feature extractor | VGG16 (torchvision) | Proven performance on image understanding tasks |
| Fusion method | Add (not concatenate) | Fewer parameters, comparable expressiveness |
| Activation | LeakyReLU (α=0.1) | Avoids dying ReLU, better gradient flow |
| Loss | CrossEntropyLoss (ignore pad) | Standard for next-token prediction |
| Teacher forcing | Full sequence at once | More efficient than prefix-pair generation |
The model is evaluated on a held-out test set (5% of Flickr8k) using Bilingual Evaluation Understudy (BLEU) scores:
| Metric | Score |
|---|---|
| BLEU-1 | ~0.55–0.62 |
| BLEU-2 | ~0.35–0.42 |
Scores depend on training duration (64 epochs) and hyperparameter tuning.
- Loss Convergence: Training loss steadily decreases over 64 epochs
- Framework: PyTorch 2.1+ with manual training loop and gradient clipping
- Hardware: ~45 min on NVIDIA Tesla T4 (Kaggle)
- Dataset: 7,600 training images × 5 captions = 38,000 caption samples
- Python 3.8+
- uv (recommended) or pip
- NVIDIA GPU recommended for training (~45 min on T4)
- Flickr8k Dataset from Kaggle
# Clone the repository
git clone https://github.com/yoel3imari/image-captioning.git
cd image-captioning
# Create a virtual environment and install all dependencies
uv sync
# Activate the virtual environment
source .venv/bin/activate # On Windows: .venv\Scripts\activate# Clone the repository
git clone https://github.com/yoel3imari/image-captioning.git
cd image-captioning
# Create and activate virtual environment
python3 -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# Install dependencies from pyproject.toml
pip install .Run the download script to fetch Flickr8k automatically:
# Install kagglehub (required for download)
uv add kagglehub
# Download Flickr8k (~1 GB) to ./data/flickr8k/
uv run python download_data.py --dataset flickr8k
# Or for Flickr30k (~6 GB):
uv run python download_data.py --dataset flickr30kThe script uses kagglehub — no Kaggle API key required for public datasets. It will:
- Download the dataset to a cache directory via
kagglehub - Copy the files into
./data/flickr8k/(or./data/flickr30k/) - Print the expected
DATA_DIRpath
Then run with the downloaded data:
export DATA_DIR="$(pwd)/data/flickr8k"
uv run jupyter notebook image-captioning-vgg16.ipynbThe dataset is expected in this structure:
/path/to/flickr8k/
├── captions.txt
└── Images/
├── 1000268201_693b08cb0e.jpg
├── 1001773457_577c3a7d70.jpg
└── ...
# Launch Jupyter
jupyter notebook image-captioning-vgg16.ipynb
# Or with JupyterLab
jupyter lab image-captioning-vgg16.ipynbRun all cells in order. The notebook will:
- Extract VGG16 features for all images (cached to disk)
- Preprocess and tokenize captions
- Train the encoder-decoder model (64 epochs)
- Evaluate with BLEU scores
- Visualize predictions on test images
image-captioning/
├── LICENSE # MIT License
├── README.md # This file
├── pyproject.toml # Project config + dependencies (uv/pip)
├── .python-version # Python version pin
├── .gitignore # Git ignore rules
├── download_data.py # Dataset download script
├── image-captioning-vgg16.ipynb # Main implementation notebook
└── data/ # Downloaded datasets (gitignored)
└── flickr8k/
├── captions.txt
└── Images/
- Dataset size: 8K images is modest — would benefit from Flickr30k or MS-COCO
- BLEU score limitations: N-gram overlap doesn't capture semantic quality
- Fixed vocabulary: Out-of-vocabulary words at inference are discarded
- No attention mechanism: Current version uses simple feature vector; spatial attention could improve fine-grained captioning
- Add Bahdanau/Luong attention — attend to spatial regions of the feature map
- Beam search decoding — replace greedy search for higher-quality captions
- Transformer decoder — replace LSTM with transformer for parallel training
- CIDEr / SPICE evaluation — additional metrics that correlate better with human judgment
- Web demo — Gradio or Streamlit interface for interactive image captioning
This project is open source under the MIT License — see the LICENSE file for details.
Built with ❤️ using TensorFlow & Keras
@yoel3imari