--- license: mit library_name: pytorch pipeline_tag: image-to-text language: - en tags: - image-captioning - image-to-text - pytorch - efficientnet - transformer-decoder - coco - computer-vision metrics: - bleu --- # Image Captioning Model This repository contains a custom PyTorch image captioning model. The model receives an input image and generates a natural-language caption describing the image. `git hub repo`: [github](https://github.com/AliSedghiye/Image_captioning) The architecture is built from two main components: 1. **Image Encoder**: EfficientNet-V2-S backbone pretrained on ImageNet. 2. **Text Decoder**: Transformer decoder that generates captions token by token. The model was trained for image caption generation using COCO-style image-caption pairs. ## Model Architecture The model follows an encoder-decoder structure: ```text Input Image ↓ EfficientNet-V2-S Image Encoder ↓ Image Feature Tokens ↓ Transformer Text Decoder ↓ Generated Caption ``` ### Image Encoder The encoder uses `EfficientNet_V2_S` from `torchvision.models`. The image encoder extracts visual features from the input image and projects them into a 256-dimensional embedding space. The final image representation is treated as a sequence of visual tokens. Encoder details: ```text Backbone: EfficientNet-V2-S Input image size: 224 x 224 Output visual tokens: 49 Embedding dimension: 256 ImageNet normalization: Yes ``` ### Text Decoder The decoder is a Transformer decoder that generates captions autoregressively. Decoder details: ```text Vocabulary size: 9,721 Embedding dimension: 256 Number of Transformer decoder layers: 6 Number of attention heads: 8 Feed-forward dimension: 1024 Maximum caption length: 52 Dropout: 0.1 Decoding methods: Greedy search and beam search ``` ## Repository Files This repository contains: ```text best_phase1.pt # PyTorch checkpoint Training-5k.ipynb # Training and inference notebook ``` The checkpoint contains: ```text epoch model val_loss ``` Checkpoint information: ```text Checkpoint file: best_phase1.pt Epoch: 8 Validation loss: 3.6158 ``` ## Important Note About Vocabulary This model uses a custom word-level vocabulary built from the training captions. The checkpoint stores the model weights, but it does **not** store the vocabulary mapping. To run inference correctly, you must use the same vocabulary that was used during training. The vocabulary contains 9,721 tokens and uses the following special tokens: ```text = 0 = 1 = 2 = 3 ``` If you want to make this model easier to use, it is recommended to upload an additional file such as: ```text vocab.json ``` containing the `stoi` and `itos` mappings. ## Training Data The model was trained using COCO-style image-caption data. The training notebook is configured to use: ```text Dataset format: COCO captions Training annotations: captions_train2014.json Validation annotations: captions_val2014.json Image size: 224 x 224 Batch size: 32 Maximum caption length: 52 ``` The notebook version included in this repository was designed for a smaller training experiment using a limited number of samples. ## Image Preprocessing Images are resized to `224 x 224` and normalized with ImageNet statistics: ```python IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] ``` Validation and inference transforms: ```python import torchvision.transforms as T transform = T.Compose([ T.Resize((224, 224)), T.ToTensor(), T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ), ]) ``` ## How to Use This is a custom PyTorch model. It is **not** a standard Hugging Face Transformers model, so it cannot be loaded directly with: ```python AutoModel.from_pretrained(...) ``` To use the model, open and run the notebook: ```text Training-5k.ipynb ``` The notebook contains: ```text Vocabulary class Dataset class EfficientNet encoder Transformer decoder ImageCaptioningModel class Training loop Checkpoint loading Greedy decoding Beam-search decoding Evaluation code ``` ## Loading the Checkpoint After defining the model architecture and rebuilding/loading the same vocabulary, the checkpoint can be loaded as follows: ```python import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ImageCaptioningModel( vocab_size=9721, embed_dim=256, num_heads=8, num_layers=6, ff_dim=1024, max_len=52, dropout=0.1 ).to(device) checkpoint = torch.load("best_phase1.pt", map_location=device) model.load_state_dict(checkpoint["model"]) model.eval() print("Loaded checkpoint") print("Epoch:", checkpoint["epoch"]) print("Validation loss:", checkpoint["val_loss"]) ``` ## Generating a Caption The notebook includes two caption generation methods: ```python model.generate_greedy(image_tensor) model.generate_beam(image_tensor, beam_size=5) ``` Example: ```python from PIL import Image image = Image.open("example.jpg").convert("RGB") image_tensor = transform(image) caption = model.generate_beam(image_tensor, beam_size=5) print(caption) ``` ## Example Output Example caption format: ```text a bicycle with a clock as the front wheel ``` Actual output quality depends on the training data size, checkpoint version, and decoding method. ## Evaluation The notebook includes BLEU evaluation code using NLTK: ```python from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction ``` You can evaluate the model on validation images using greedy decoding or beam search. Recommended metrics for this task: ```text BLEU-1 BLEU-4 CIDEr METEOR ROUGE-L ``` ## Limitations This model is an experimental image captioning model. Known limitations: * The model uses a custom word-level tokenizer, not a subword tokenizer. * The vocabulary must match the original training vocabulary. * The checkpoint alone is not enough for fully reproducible inference unless the vocabulary is also available. * Caption quality may be limited if the model was trained on a small subset of the dataset. * The model may generate generic or repetitive captions. * The model may fail on images that are very different from the training distribution. * The model may hallucinate objects that are not present in the image. ## Recommended Improvements To make this repository easier to use, future versions should include: ```text vocab.json model.py requirements.txt inference.py example images evaluation results ``` A better repository structure would be: ```text . ├── README.md ├── best_phase1.pt ├── Training-5k.ipynb ├── vocab.json ├── model.py ├── inference.py └── requirements.txt ``` ## Requirements The notebook uses the following main libraries: ```text torch torchvision Pillow numpy matplotlib nltk pycocotools pycocoevalcap einops ``` Install dependencies with: ```bash pip install torch torchvision pillow numpy matplotlib nltk pycocotools pycocoevalcap einops ``` ## Citation If you use this model, please cite or mention this repository. ## Author Created as a custom PyTorch image captioning model using an EfficientNet image encoder and a Transformer text decoder.