OpenVision 2 — Image Captioning (Unofficial Reimplementation)

Course project for CAP6415: Computer Vision (Fall 2025). An educational, from‑scratch reimplementation of the OpenVision 2 idea: caption images by pairing a frozen Vision Transformer with a lightweight GPT‑style text decoder that cross‑attends to visual patch tokens.

This is not the official OpenVision 2 code. It is a compact, readable implementation built for learning, trained on a single GPU (Apple Silicon / MPS) on Flickr8k.

---

Table of Contents

• What it does
• Architecture
• Results
• Repository layout
• Installation
• Dataset setup (Flickr8k)
• Training
• Inference
• Demo
• How it works
• Limitations & future work
• Acknowledgments

---

What it does

Given an image, the model generates a natural‑language caption:

Input	Generated caption (epoch 10)
Girl on stairs	"a little girl in a pink dress is standing on the stairs"
Two dogs	"two dogs are playing in the grass"
Girl portrait	"a young girl wearing a hat is smiling"

The pipeline is a classic encoder–decoder:

Image ──▶ ViT encoder ──▶ projection ──▶ Transformer decoder ──▶ Caption
        (frozen, 197×768)   (768→512)      (6 layers, cross-attn)

---

Architecture

Component	Details
Visual encoder	vit_base_patch16_224 from timm, pretrained, frozen. Classification head removed; forward_features returns 197 tokens (196 patches + CLS) of dim 768.
Projection layer	Single nn.Linear(768 → 512), trainable. Maps visual features into the decoder's embedding space.
Causal decoder	nn.TransformerDecoder, 6 layers, 8 heads, d_model=512, FFN 2048, GELU, dropout 0.1. Self‑attention is causally masked; cross‑attention uses the visual tokens as memory.
Visual token masking	During training, 50% of patch tokens are randomly zeroed (CLS kept). Disabled at validation / inference.
Tokenizer	SimpleTokenizer — word‑level, lowercase, punctuation stripped. Special tokens: <PAD>=0, <START>=1, <END>=2, <UNK>=3. Vocab size 5,000, min_freq=2.
Objective	Causal language modeling (next‑token cross‑entropy), teacher forcing.
Decoding	Greedy / temperature sampling, stops at <END> or max_length.

Model size: ~116.6M parameters total, ~30.8M trainable (projection + decoder); the ViT encoder (~85.8M) is frozen.

---

Results

Trained on Flickr8k for 10 epochs (batch size 32, Adam lr=1e-4, ~30 min/epoch on MPS):

Metric	Value
Best validation loss	~0.89
Final training loss	~0.41
Vocabulary size	5,000

Caption quality improves steadily: epoch‑1 outputs are repetitive and full of <UNK>; by epoch ~7–10 they are grammatical and contextually relevant. No BLEU/CIDEr/METEOR scoring is implemented yet (see future work).

Numbers are from the training runs recorded in development_logs/.

---

Repository layout

.
├── src/
│   ├── model.py            # ViT encoder, projection, causal decoder, full model + loss
│   ├── tokenizer.py        # SimpleTokenizer (word-level, JSON save/load)
│   │
│   ├── real_dataset.py     # Flickr8k Dataset + dataloader factory (split by image)
│   ├── train_real.py       # Training on Flickr8k  ── main entry point
│   ├── inference_real.py   # Caption generation from a trained checkpoint
│   │
│   ├── dataset.py          # Prototype: synthetic random data (no download needed)
│   ├── train.py            # Prototype: training loop on synthetic data
│   └── inference.py        # Prototype: decoding utilities
│
├── demo_6_images.py        # Generate captions for 6 hand-picked demo images
├── demo_6_images.sh        # Shell wrapper for the demo
├── data/
│   └── captions.txt        # Flickr8k captions (images downloaded separately)
├── checkpoints/            # Trained model + tokenizer (model .pt files via Git LFS)
├── development_logs/        # Weekly progress logs (week1–week4)
└── requirements.txt

Two tracks: *_real.py files are the real pipeline (Flickr8k). dataset.py / train.py / inference.py are an earlier prototype that runs on randomly generated tensors — useful for smoke‑testing the architecture with no dataset.

---

Installation

Git LFS is required. The trained checkpoint in checkpoints/ (~712 MB) is stored via Git LFS. Install it before cloning, or the .pt files will only be pointer stubs.

# one-time
git lfs install

git clone https://github.com/v6y4k2s/CAP6415_F25_project_OpenVision2_Code_implementation.git
cd CAP6415_F25_project_OpenVision2_Code_implementation

python3 -m venv venv
source venv/bin/activate        # Windows: venv\Scripts\activate
pip install -r requirements.txt

Requirements: torch, torchvision, timm, numpy, tqdm, matplotlib, Pillow. Runs on CUDA, Apple Silicon (MPS), or CPU — the scripts auto‑detect.

---

Dataset setup (Flickr8k)

The images are not included. Download them from Kaggle: https://www.kaggle.com/datasets/adityajn105/flickr8k

Arrange the files so the layout is:

data/
├── Images/          # ~8,091 .jpg files
│   ├── 1000268201_693b08cb0e.jpg
│   └── ...
└── captions.txt     # already in the repo (image,caption CSV; 5 captions/image)

The loader splits by image (80% train / 10% val / 10% test) so all captions for one image stay in the same split.

---

Training

python src/train_real.py --epochs 10 --batch-size 32 --vocab-size 5000 --save-model

Common flags (see python src/train_real.py --help for all):

Flag	Default	Meaning
--epochs	10	training epochs
--batch-size	32	batch size
--vocab-size	5000	max vocabulary
--min-freq	2	drop words rarer than this
--max-caption-length	50	pad/truncate length
--text-dim	512	decoder embedding dim
--num-decoder-layers	6	decoder depth
--num-heads	8	attention heads
--lr	1e-4	Adam learning rate
--save-model	off	save best checkpoint + tokenizer
--output-dir	checkpoints	where to write best_model_real.pt and tokenizer.json

The best model (lowest val loss) and its tokenizer are saved together so inference is fully reproducible.

---

Inference

python src/inference_real.py \
  --checkpoint checkpoints/best_model_real.pt \
  --tokenizer  checkpoints/tokenizer.json \
  --image      path/to/image.jpg

Options: --max-length (default 50), --temperature (1.0 = greedy), --show-variants (print captions at temperatures 0.7 / 1.0 / 1.3).

---

Demo

Generate captions for 6 diverse sample images (people, animals, actions):

python demo_6_images.py \
  --checkpoint checkpoints/best_model_real.pt \
  --tokenizer  checkpoints/tokenizer.json

# or
bash demo_6_images.sh

---

How it works

1. Encode the image with a frozen ViT‑B/16 → (B, 197, 768) patch tokens.
2. Project each token 768 → 512 so it lives in the decoder's space.
3. Mask ~50% of patch tokens at random during training (robustness / augmentation).
4. Decode the caption autoregressively: the decoder self‑attends over previously generated words (causal mask) and cross‑attends over the visual tokens as memory.
5. Loss: shift the caption by one and apply cross‑entropy for next‑token prediction (caption[:, :-1] → caption[:, 1:]).
6. Generate: start from <START>, greedily append the arg‑max token, stop at <END>.

---

Limitations & future work

• Word‑level vocabulary (no BPE / subword) → <UNK> on rare words.
• Greedy decoding only — no beam search.
• No caption metrics — BLEU / CIDEr / METEOR not implemented.
• Encoder frozen — ViT is never fine‑tuned.
• Small dataset — Flickr8k only.

Planned: beam search, BLEU/CIDEr evaluation, attention‑map visualization, ViT fine‑tuning, and MS‑COCO support.

---

Acknowledgments

• Vision Transformer weights via timm
• Flickr8k dataset
• Built on the PyTorch ecosystem
• Concept inspired by the OpenVision 2 vision‑language architecture (this repo is an unofficial reimplementation)