Image Captioning System that generates natural language captions for any image.
The architecture for the model is inspired from "Show and Tell" [1] by Vinyals et al. The model is built using Keras library.
The project also contains code for Attention LSTM layer, although not integrated in the model.
The model is trained on Flickr8k Dataset
Although it can be trained on others like Flickr30k or MS COCO
The model has been trained for 20 epoches on 6000 training samples of Flickr8k Dataset. It acheives a BLEU-1 = ~0.59 with 1000 testing samples.
- tensorflow
- keras
- numpy
- h5py
- progressbar2
These requirements can be easily installed by:
pip install -r requirements.txt
- caption_generator.py: The base script that contains functions for model creation, batch data generator etc.
- prepare_data.py: Extracts features from images using VGG16 imagenet model. Also prepares annotation for training. Changes have to be done to this script if new dataset is to be used.
- train_model.py: Module for training the caption generator.
- eval_model.py: Contains module for evaluating and testing the performance of the caption generator, currently, it contains the BLEU metric.
- Download pre-trained weights from releases
- Move
model_weight.h5tomodelsdirectory - Prepare data using
python prepare_data.py - For inference on example image, run:
python eval_model.py -i [img-path]
After the requirements have been installed, the process from training to testing is fairly easy. The commands to run:
python prepare_data.pypython train_model.pypython eval_model.py
After training, evaluation on an example image can be done by running:
python eval_model.py -m [model-checkpoint] -i [img-path]
| Image | Caption |
|---|---|
 |
Generated Caption: A white and black dog is running through the water |
 |
Generated Caption: man is skiing on snowy hill |
 |
Generated Caption: man in red shirt is walking down the street |
[1] Oriol Vinyals, Alexander Toshev, Samy Bengio, Dumitru Erhan. Show and Tell: A Neural Image Caption Generator
[2] Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhutdinov, Richard Zemel, Yoshua Bengio. Show, Attend and Tell: Neural Image Caption Generation with Visual Attention
MIT License. See LICENSE file for details.