""" Model definitions for the Khmer image captioning model (frozen ResNet-101 encoder + Bahdanau-attention LSTM decoder, trained from scratch). Usage: from modeling_khmer_captioning import load_model, caption_image encoder, decoder, itos, stoi, cfg = load_model(".") # dir with config.json/decoder.pt/vocab.json caption = caption_image("photo.jpg", encoder, decoder, itos, stoi, cfg) """ import json import os 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 PAD_TOKEN, START_TOKEN, END_TOKEN, UNK_TOKEN = "", "", "", "" class FrozenCNNEncoder(nn.Module): def __init__(self, encoded_image_size=14): 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=256, attention_dim=256, decoder_dim=512, encoder_dim=2048, 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 def load_model(repo_dir, device="cpu"): with open(os.path.join(repo_dir, "config.json"), "r", encoding="utf-8") as f: cfg = json.load(f) with open(os.path.join(repo_dir, "vocab.json"), "r", encoding="utf-8") as f: vocab_data = json.load(f) itos = vocab_data["itos"] stoi = {w: i for i, w in enumerate(itos)} encoder = FrozenCNNEncoder(encoded_image_size=cfg["encoded_image_size"]).to(device) encoder.eval() decoder = AttentionDecoder( vocab_size=cfg["vocab_size"], embed_dim=cfg["embed_dim"], attention_dim=cfg["attention_dim"], decoder_dim=cfg["decoder_dim"], encoder_dim=cfg["encoder_dim"], dropout=cfg["dropout"], ).to(device) state_dict = torch.load(os.path.join(repo_dir, "decoder.pt"), map_location=device) decoder.load_state_dict(state_dict) decoder.eval() return encoder, decoder, itos, stoi, cfg def _image_transform(image_size): return 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]), ]) def generate_caption(image_tensor, encoder, decoder, itos, stoi, cfg, device="cpu", beam_size=3): vocab_size = cfg["vocab_size"] encoder_dim = cfg["encoder_dim"] max_len = cfg.get("max_caption_len", 400) 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.squeeze(1).tolist() best_seq = complete_seqs[int(np.argmax(complete_scores))] words = [itos[idx] for idx in best_seq if idx not in (stoi[START_TOKEN], stoi[END_TOKEN], stoi[PAD_TOKEN])] return "".join(words) def caption_image(image_path, encoder, decoder, itos, stoi, cfg, device="cpu", beam_size=3): image = Image.open(image_path) if image.mode != "RGB": image = image.convert("RGB") tensor = _image_transform(cfg.get("image_size", 224))(image) return generate_caption(tensor, encoder, decoder, itos, stoi, cfg, device=device, beam_size=beam_size)