--- language: km tags: - image-captioning - khmer - resnet - transformer - pytorch license: apache-2.0 --- # 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](https://huggingface.co/datasets/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 ```python !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")) ```