| """ |
| 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 = "<pad>", "<start>", "<end>", "<unk>" |
|
|
|
|
| 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) |
|
|
