phonsobon/khmer_images_captioning_v3

Khmer Image Captioning (CNN Encoder + Attention LSTM Decoder)

Trained from scratch on phonsobon/khmer_images_captioning_v2.

Model repo: phonsobon/khmer_images_captioning_v3

• Encoder: frozen ResNet-101 (ImageNet weights, feature extraction only)
• Decoder: LSTM with Bahdanau attention, trained from scratch

Files

• decoder_best.pt: decoder weights
• vocab.json: word-to-index / index-to-word mappings
• config.json: model hyperparameters

How to test / run inference

import json
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image
from huggingface_hub import hf_hub_download

REPO_ID = "phonsobon/khmer_images_captioning_v3"
device = "cuda" if torch.cuda.is_available() else "cpu"

PAD_TOKEN, START_TOKEN, END_TOKEN, UNK_TOKEN = "<pad>", "<start>", "<end>", "<unk>"

# 1. Download files from the Hub
decoder_path = hf_hub_download(REPO_ID, "decoder_best.pt")
config_path = hf_hub_download(REPO_ID, "config.json")
vocab_json_path = hf_hub_download(REPO_ID, "vocab.json")

# 2. Load vocab and config
with open(vocab_json_path, "r", encoding="utf-8") as f:
    vocab = json.load(f)
itos, stoi = vocab["itos"], vocab["stoi"]

with open(config_path, "r", encoding="utf-8") as f:
    config = json.load(f)

VOCAB_SIZE = config["vocab_size"]
ENCODER_DIM = config["encoder_dim"]

# 3. Model definitions

class FrozenCNNEncoder(nn.Module):
    def __init__(self, encoded_image_size=config["encoded_image_size"]):
        super().__init__()
        resnet = models.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V2)
        modules = list(resnet.children())[:-2]
        self.resnet = nn.Sequential(*modules)
        self.adaptive_pool = nn.AdaptiveAvgPool2d((encoded_image_size, encoded_image_size))
        for param in self.resnet.parameters():
            param.requires_grad = False
        self.resnet.eval()

    @torch.no_grad()
    def forward(self, images):
        features = self.resnet(images)
        features = self.adaptive_pool(features)
        features = features.permute(0, 2, 3, 1)
        return features


class Attention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        self.encoder_att = nn.Linear(encoder_dim, attention_dim)
        self.decoder_att = nn.Linear(decoder_dim, attention_dim)
        self.full_att = nn.Linear(attention_dim, 1)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, encoder_out, decoder_hidden):
        att1 = self.encoder_att(encoder_out)
        att2 = self.decoder_att(decoder_hidden).unsqueeze(1)
        att = self.full_att(self.relu(att1 + att2)).squeeze(2)
        alpha = self.softmax(att)
        context = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)
        return context, alpha


class AttentionDecoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, attention_dim, decoder_dim, encoder_dim, dropout=0.5):
        super().__init__()
        self.vocab_size = vocab_size
        self.decoder_dim = decoder_dim
        self.attention = Attention(encoder_dim, decoder_dim, attention_dim)
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.dropout = nn.Dropout(dropout)
        self.lstm_cell = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)
        self.init_h = nn.Linear(encoder_dim, decoder_dim)
        self.init_c = nn.Linear(encoder_dim, decoder_dim)
        self.f_beta = nn.Linear(decoder_dim, encoder_dim)
        self.sigmoid = nn.Sigmoid()
        self.fc = nn.Linear(decoder_dim, vocab_size)

    def init_hidden_state(self, encoder_out):
        mean_encoder_out = encoder_out.mean(dim=1)
        h = self.init_h(mean_encoder_out)
        c = self.init_c(mean_encoder_out)
        return h, c


# 4. Rebuild encoder + decoder and load weights
encoder = FrozenCNNEncoder(encoded_image_size=config["encoded_image_size"]).to(device).eval()
decoder = AttentionDecoder(
    vocab_size=config["vocab_size"],
    embed_dim=config["embed_dim"],
    attention_dim=config["attention_dim"],
    decoder_dim=config["decoder_dim"],
    encoder_dim=config["encoder_dim"],
).to(device)
checkpoint = torch.load(decoder_path, map_location=device)
decoder.load_state_dict(checkpoint["model_state_dict"])
decoder.eval()

print(f"Loaded checkpoint from epoch {checkpoint['epoch']}, "
      f"best val loss: {checkpoint['best_val_loss']:.4f}")

# 5. Beam-search caption generation
def generate_caption(image_tensor, beam_size=3, max_len=50):
    decoder.eval()
    with torch.no_grad():
        image_tensor = image_tensor.unsqueeze(0).to(device)
        encoder_out = encoder(image_tensor)
        encoder_out = encoder_out.view(1, -1, encoder_out.size(-1))
        num_regions = encoder_out.size(1)
        encoder_out = encoder_out.expand(beam_size, num_regions, ENCODER_DIM)
        k = beam_size
        seqs = torch.full((k, 1), stoi[START_TOKEN], dtype=torch.long, device=device)
        top_k_scores = torch.zeros(k, 1, device=device)
        h, c = decoder.init_hidden_state(encoder_out)
        complete_seqs, complete_scores = [], []
        step = 1
        while True:
            embeddings = decoder.embedding(seqs[:, -1])
            context, _ = decoder.attention(encoder_out, h)
            gate = decoder.sigmoid(decoder.f_beta(h))
            context = gate * context
            h, c = decoder.lstm_cell(torch.cat([embeddings, context], dim=1), (h, c))
            scores = F.log_softmax(decoder.fc(h), dim=1)
            scores = top_k_scores.expand_as(scores) + scores
            if step == 1:
                top_k_scores, top_k_words = scores[0].topk(k, dim=0)
            else:
                top_k_scores, top_k_words = scores.view(-1).topk(k, dim=0)
            prev_seq_inds = top_k_words // VOCAB_SIZE
            next_word_inds = top_k_words % VOCAB_SIZE
            seqs = torch.cat([seqs[prev_seq_inds], next_word_inds.unsqueeze(1)], dim=1)
            incomplete = [i for i, w in enumerate(next_word_inds) if w.item() != stoi[END_TOKEN]]
            complete = [i for i in range(len(next_word_inds)) if i not in incomplete]
            if complete:
                complete_seqs.extend(seqs[complete].tolist())
                complete_scores.extend(top_k_scores[complete].tolist())
            k -= len(complete)
            if k == 0 or step >= max_len:
                break
            seqs = seqs[incomplete]
            h, c = h[prev_seq_inds][incomplete], c[prev_seq_inds][incomplete]
            encoder_out = encoder_out[prev_seq_inds][incomplete]
            top_k_scores = top_k_scores[incomplete].unsqueeze(1)
            step += 1
        if not complete_seqs:
            complete_seqs, complete_scores = seqs.tolist(), top_k_scores.view(-1).tolist()
        best_seq = complete_seqs[int(np.argmax(complete_scores))]
        words = [itos[str(idx)] if str(idx) in itos else itos[idx]
                 for idx in best_seq
                 if idx not in (stoi[START_TOKEN], stoi[END_TOKEN], stoi[PAD_TOKEN])]
        return "".join(words)


# 6. Preprocess a test image and generate a caption
transform = transforms.Compose([
    transforms.Resize((config["image_size"], config["image_size"])),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
image = Image.open("test.jpg").convert("RGB")
image_tensor = transform(image)

caption = generate_caption(image_tensor, beam_size=3)
print("Predicted caption:", caption)