khmer-image-captioning-cnn-transformer

Khmer Image Captioning (CNN Encoder + Transformer Decoder)

ResNet-50 encoder (ImageNet-pretrained, layer4 fine-tuned) + Transformer decoder (multi-head self-attention, sinusoidal positional encoding), trained on phonsobon/khmer_images_captioning_v2 to generate Khmer-language image captions.

Tokenizer: khmer-nltk word segmentation followed by SentencePiece BPE (vocab size VOCAB_SIZE) trained on the segmented corpus. Khmer script has no spaces between words, so captions are first segmented into real words with khmer-nltk's CRF-based tokenizer, then BPE is trained on top (constrained to not merge across word boundaries), giving linguistically-aware subword units instead of purely frequency-driven character fragments.

This is a custom PyTorch architecture (not a standard transformers model class), so loading it requires re-declaring the model classes below, then loading the weights from pytorch_model.pt. The snippet below is self-contained: it downloads everything from this repo and runs a test caption end to end.

How to use

!pip install -q huggingface_hub sentencepiece khmer-nltk torch torchvision pillow
 
import json, math
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from torchvision.models import resnet50
from huggingface_hub import hf_hub_download
import sentencepiece as spm
 
REPO_ID = "__REPO_ID__"
 
# ---- download the model files from this repo ----
config_path = hf_hub_download(REPO_ID, "config.json")
weights_path = hf_hub_download(REPO_ID, "pytorch_model.pt")
tokenizer_path = hf_hub_download(REPO_ID, "khmer_bpe.model")
 
with open(config_path, "r", encoding="utf-8") as f:
    config = json.load(f)
 
sp = spm.SentencePieceProcessor(model_file=tokenizer_path)
PAD_ID, UNK_ID, BOS_ID, EOS_ID = config["pad_id"], config["unk_id"], config["bos_id"], config["eos_id"]
IMAGE_SIZE, MAX_LEN, D_MODEL = config["image_size"], config["max_len"], config["d_model"]
 
# ---- model definition (must match training architecture) ----
class CNNEncoder(nn.Module):
    def __init__(self, d_model=512):
        super().__init__()
        resnet = resnet50(weights=None)  # weights are loaded from pytorch_model.pt
        self.backbone = nn.Sequential(*list(resnet.children())[:-2])
        self.proj = nn.Conv2d(2048, d_model, kernel_size=1)
        self.num_patches = (IMAGE_SIZE // 32) ** 2
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, d_model))
 
    def forward(self, images):
        feats = self.backbone(images)
        feats = self.proj(feats)
        B, C, H, W = feats.shape
        feats = feats.flatten(2).transpose(1, 2)
        return feats + self.pos_embed
 
 
class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=200):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))
 
    def forward(self, x):
        return x + self.pe[:, : x.size(1)]
 
 
class TransformerCaptionDecoder(nn.Module):
    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dim_ff=2048, dropout=0.1):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
        self.pos_enc = SinusoidalPositionalEncoding(d_model)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, dropout=dropout, batch_first=True,
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)
        self.d_model = d_model
 
    def forward(self, tgt_in, memory, tgt_key_padding_mask=None):
        L = tgt_in.size(1)
        causal_mask = nn.Transformer.generate_square_subsequent_mask(L).to(tgt_in.device)
        x = self.embed(tgt_in) * math.sqrt(self.d_model)
        x = self.pos_enc(x)
        out = self.decoder(tgt=x, memory=memory, tgt_mask=causal_mask, tgt_key_padding_mask=tgt_key_padding_mask)
        return self.fc_out(out)
 
 
class ImageCaptioningModel(nn.Module):
    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dim_ff=2048, dropout=0.1):
        super().__init__()
        self.encoder = CNNEncoder(d_model=d_model)
        self.decoder = TransformerCaptionDecoder(vocab_size, d_model, nhead, num_layers, dim_ff, dropout)
 
    def forward(self, images, tgt_in, tgt_key_padding_mask=None):
        memory = self.encoder(images)
        return self.decoder(tgt_in, memory, tgt_key_padding_mask=tgt_key_padding_mask)
 
 
# ---- load weights ----
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ImageCaptioningModel(
    vocab_size=config["vocab_size"], d_model=D_MODEL,
    nhead=config["nhead"], num_layers=config["num_decoder_layers"], dim_ff=config["dim_feedforward"],
).to(device)
model.load_state_dict(torch.load(weights_path, map_location=device))
model.eval()
 
# ---- inference on a test image ----
tfms = transforms.Compose([
    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
 
@torch.no_grad()
def generate_caption(image_path, max_len=MAX_LEN):
    image = Image.open(image_path).convert("RGB")
    image_tensor = tfms(image).unsqueeze(0).to(device)
    memory = model.encoder(image_tensor)
 
    ids = [BOS_ID]
    for _ in range(max_len):
        tgt = torch.tensor([ids], device=device)
        logits = model.decoder(tgt, memory)
        next_id = logits[0, -1].argmax().item()
        if next_id == EOS_ID:
            break
        ids.append(next_id)
    return sp.decode(ids[1:])
 
print(generate_caption("your_test_image.jpg"))