This project is a transformer-based image captioning system that accepts an image as input and generates a descriptive caption.
The model pairs a Vision Transformer (ViT) encoder with a GPT-2 style decoder, consisting of 12 ViT encoder layers and 12 GPT-2 decoder layers.
Training was performed on the Flickr30k dataset, which contains 30,000 images with 5 captions per image. The model was trained for 8 epochs, taking approximately 8 hours on 2 Kaggle T4 GPUs.
The model achieved the following evaluation scores:
BLEU-4 0.3034
ROUGE-L 0.4930
METEOR 0.4937
CIDEr 0.8220
Below are the loss curves:
To generate captions for your own images, download models.py, app.py, and requirements.txt, then run:
pip install -r requirements.txt
streamlit run app.py
Note: the first prediction run will take longer because the model weights (~1GB) must be downloaded. This is what appears when you first run the application:
Below are a few example predictions on the validation dataset:
These example predictions are a starting point; quality can be improved by further training and/or using a larger dataset such as MS COCO (≈500k captioned images).
See the full notebook or the Kaggle version.
Download the pretrained model weights.