Spaces:
Sleeping
Sleeping
| import torch | |
| import pickle | |
| import string | |
| import torch.nn.functional as F | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| def decoder(indices): | |
| with open("vocab.pkl", 'rb') as f: | |
| vocab = pickle.load(f) | |
| tokens = [vocab.lookup_token(idx) for idx in indices] | |
| words = [] | |
| current_word = [] | |
| for token in tokens: | |
| if len(token) == 1 and token in string.ascii_lowercase: | |
| current_word.append(token) | |
| else: | |
| if current_word: | |
| words.append("".join(current_word)) | |
| current_word = [] | |
| words.append(token) | |
| if current_word: | |
| words.append(" "+"".join(current_word)) | |
| return "".join(words) | |
| def beam_search_caption(model, images, vocab, decoder, device="cpu", | |
| start_token="<sos>", end_token="<eos>", | |
| beam_width=3, max_seq_length=100): | |
| model.eval() | |
| with torch.no_grad(): | |
| start_index = vocab[start_token] | |
| end_index = vocab[end_token] | |
| images = images.to(device) | |
| batch_size = images.size(0) | |
| # Ensure batch_size is 1 for beam search (one image at a time) | |
| if batch_size != 1: | |
| raise ValueError("Beam search currently supports batch_size=1.") | |
| cnn_features = model.cnn(images) # (B, 49, 2048) | |
| h, c = model.lstm.init_hidden_state(batch_size) | |
| word_input = torch.full((batch_size,), start_index, dtype=torch.long).to(device) | |
| embeddings = model.lstm.embedding(word_input) | |
| context, _ = model.lstm.attention(cnn_features, h[-1]) | |
| lstm_input = torch.cat([embeddings, context], dim=1).unsqueeze(1) | |
| sequences = [([start_index], 0.0, lstm_input, (h, c))] # List of tuples: (sequence, score, input, state) | |
| completed_sequences = [] | |
| for _ in range(max_seq_length): | |
| all_candidates = [] | |
| for seq, score, lstm_input, (h,c) in sequences: | |
| if seq[-1] == end_index: | |
| completed_sequences.append((seq, score)) | |
| continue | |
| lstm_out, (h_new, c_new) = model.lstm.lstm(lstm_input, (h, c)) # lstm_out: (1, 1, 1024) | |
| output = model.lstm.fc(lstm_out.squeeze(1)) # Shape: (1, vocab_size) | |
| log_probs = F.log_softmax(output, dim=1) # Shape: (1, vocab_size) | |
| top_log_probs, top_indices = log_probs.topk(beam_width, dim=1) # Each of shape: (1, beam_width) | |
| for i in range(beam_width): | |
| token = top_indices[0, i].item() | |
| token_log_prob = top_log_probs[0, i].item() | |
| new_seq = seq + [token] | |
| new_score = score + token_log_prob | |
| token_tensor = torch.tensor([token], device=device) | |
| embeddings = model.lstm.embedding(token_tensor) | |
| context, _ = model.lstm.attention(cnn_features, h_new[-1]) | |
| new_lstm_input = torch.cat([embeddings, context], dim=1).unsqueeze(1) | |
| if h_new is not None and c_new is not None: | |
| h_new, c_new = (h_new.clone(), c_new.clone()) | |
| else: | |
| h_new, c_new = None, None | |
| all_candidates.append((new_seq, new_score, new_lstm_input, (h_new, c_new) )) | |
| if not all_candidates: | |
| break | |
| ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True) | |
| sequences = ordered[:beam_width] | |
| if len(completed_sequences) >= beam_width: | |
| break | |
| if len(completed_sequences) == 0: | |
| completed_sequences = sequences | |
| best_seq = max(completed_sequences, key=lambda x: x[1]) | |
| best_caption = decoder(best_seq[0]) | |
| return best_caption | |
| ## ResNet50 (CNN Encoder) | |
| class ResNet50(nn.Module): | |
| def __init__(self): | |
| super(ResNet50, self).__init__() | |
| self.ResNet50 = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) | |
| self.features = nn.Sequential(*list(self.ResNet50.children())[:-2]) | |
| self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) | |
| for param in self.ResNet50.parameters(): | |
| param.requires_grad = False | |
| def forward(self, x): | |
| x = self.features(x) | |
| x = self.avgpool(x) | |
| B, C, H, W = x.size() | |
| x = x.view(B, C, -1) # Flatten spatial dimensions: (B, 2048, 49) | |
| x = x.permute(0, 2, 1) # (B, 49, 2048) - 49 spatial locations | |
| return x | |
| class Attention(nn.Module): | |
| def __init__(self, feature_size, hidden_size): | |
| super(Attention, self).__init__() | |
| self.attention = nn.Linear(feature_size + hidden_size, hidden_size) | |
| self.attn_weights = nn.Linear(hidden_size, 1) | |
| def forward(self, features, hidden_state): # features: (B, 49, 2048), hidden_state: (B, hidden_size) | |
| hidden_state = hidden_state.unsqueeze(1).repeat(1, features.size(1), 1) # (B, 49, hidden_size) | |
| combined = torch.cat((features, hidden_state), dim=2) # (B, 49, feature_size + hidden_size) | |
| attn_hidden = torch.tanh(self.attention(combined)) # (B, 49, hidden_size) | |
| attention_logits = self.attn_weights(attn_hidden).squeeze(2) # (B, 49) | |
| attention_weights = torch.softmax(attention_logits, dim=1) # (B, 49) | |
| context = (features * attention_weights.unsqueeze(2)).sum(dim=1) # (B, 2048) | |
| return context, attention_weights | |
| # Attention without learnable paramters: | |
| # logits = torch.matmul(features, hidden_state.unsqueeze(2)) # (B, 49, 1) - Batch Matriax | |
| # attention_weights = torch.softmax(logits, dim=1).squeeze(2) # (B, 49) | |
| # context = (features * attention_weights.unsqueeze(2)).sum(dim=1) # (B, 2048) | |
| class lstm(nn.Module): | |
| def __init__(self, feature_size, hidden_size, number_layers, embedding_dim, vocab_size): | |
| super(lstm, self).__init__() | |
| self.hidden_size = hidden_size | |
| self.embedding = nn.Embedding(vocab_size, hidden_size) | |
| self.attention = Attention(feature_size, hidden_size) | |
| self.lstm = nn.LSTM( | |
| input_size=hidden_size + feature_size, # input: concatenated context and word embedding | |
| hidden_size=hidden_size, | |
| num_layers=number_layers, | |
| dropout=0.5, | |
| batch_first=True, | |
| ) | |
| self.fc = nn.Linear(hidden_size, vocab_size) | |
| def forward(self, features, captions=None, max_seq_len=None, teacher_forcing_ratio=0.90): | |
| batch_size = features.size(0) | |
| max_seq_len = max_seq_len if max_seq_len is not None else captions.size(1) | |
| h, c = self.init_hidden_state(batch_size) | |
| outputs = torch.zeros(batch_size, max_seq_len, self.fc.out_features).to(features.device) | |
| word_input = torch.tensor(2, dtype=torch.long).expand(batch_size).to(features.device) # vocab["<sos>"] ---> 2 | |
| for t in range(1, max_seq_len): | |
| embeddings = self.embedding(word_input) | |
| context, _ = self.attention(features, h[-1]) | |
| lstm_input_step = torch.cat([embeddings, context], dim=1).unsqueeze(1) # Combine context + word embedding | |
| out, (h, c) = self.lstm(lstm_input_step, (h, c)) | |
| output = self.fc(out.squeeze(1)) | |
| outputs[:, t, :] = output | |
| top1 = output.argmax(1) | |
| if captions is not None and torch.rand(1).item() < teacher_forcing_ratio: | |
| word_input = captions[:, t] | |
| else: | |
| word_input = top1 | |
| return outputs | |
| def init_hidden_state(self, batch_size): | |
| device = next(self.parameters()).device | |
| h0 = torch.zeros(self.lstm.num_layers, batch_size, self.hidden_size).to(device) | |
| c0 = torch.zeros(self.lstm.num_layers, batch_size, self.hidden_size).to(device) | |
| return (h0, c0) | |
| class ImgCap(nn.Module): | |
| def __init__(self, feature_size, lstm_hidden_size, num_layers, vocab_size, embedding_dim): | |
| super(ImgCap, self).__init__() | |
| self.cnn = ResNet50() | |
| self.lstm = lstm(feature_size, lstm_hidden_size, num_layers, embedding_dim, vocab_size) | |
| def forward(self, images, captions): | |
| cnn_features = self.cnn(images) | |
| output = self.lstm(cnn_features, captions) | |
| return output | |
| def generate_caption(self, images, vocab, decoder, device="cpu", start_token="<sos>", end_token="<eos>", max_seq_length=100): | |
| self.eval() | |
| with torch.no_grad(): | |
| start_index = vocab[start_token] | |
| end_index = vocab[end_token] | |
| images = images.to(device) | |
| batch_size = images.size(0) | |
| captions = [[start_index,] for _ in range(batch_size)] | |
| end_token_appear = [False] * batch_size | |
| cnn_features = self.cnn(images) # (B, 49, 2048) | |
| h, c = self.lstm.init_hidden_state(batch_size) | |
| word_input = torch.full((batch_size,), start_index, dtype=torch.long).to(device) | |
| for t in range(max_seq_length): | |
| embeddings = self.lstm.embedding(word_input) | |
| context, _ = self.lstm.attention(cnn_features, h[-1]) # Attention context | |
| lstm_input_step = torch.cat([embeddings, context], dim=1).unsqueeze(1) # Combine context + word embedding | |
| out, (h, c) = self.lstm.lstm(lstm_input_step, (h, c)) | |
| output = self.lstm.fc(out.squeeze(1)) # (B, vocab_size) | |
| # Get the predicted word (greedy search) | |
| predicted_word_indices = torch.argmax(output, dim=1) # (B,) | |
| word_input = predicted_word_indices | |
| for i in range(batch_size): | |
| if not end_token_appear[i]: | |
| predicted_word = vocab.lookup_token(predicted_word_indices[i].item()) | |
| if predicted_word == end_token: | |
| captions[i].append(predicted_word_indices[i].item()) | |
| end_token_appear[i] = True | |
| else: | |
| captions[i].append(predicted_word_indices[i].item()) | |
| if all(end_token_appear): # Stop if all captions have reached the <eos> token | |
| break | |
| captions = [decoder(caption) for caption in captions] | |
| return captions | |
| def beam_search_caption(self, images, vocab, decoder, device="cpu", | |
| start_token="<sos>", end_token="<eos>", | |
| beam_width=3, max_seq_length=100): | |
| self.eval() | |
| with torch.no_grad(): | |
| start_index = vocab[start_token] | |
| end_index = vocab[end_token] | |
| images = images.to(device) | |
| batch_size = images.size(0) | |
| # Ensure batch_size is 1 for beam search (one image at a time) | |
| if batch_size != 1: | |
| raise ValueError("Beam search currently supports batch_size=1.") | |
| cnn_features = self.cnn(images) # (B, 49, 2048) | |
| h, c = self.lstm.init_hidden_state(batch_size) | |
| word_input = torch.full((batch_size,), start_index, dtype=torch.long).to(device) | |
| embeddings = self.lstm.embedding(word_input) | |
| context, _ = self.lstm.attention(cnn_features, h[-1]) | |
| lstm_input = torch.cat([embeddings, context], dim=1).unsqueeze(1) | |
| sequences = [([start_index], 0.0, lstm_input, (h, c))] # List of tuples: (sequence, score, input, state) | |
| completed_sequences = [] | |
| for _ in range(max_seq_length): | |
| all_candidates = [] | |
| for seq, score, lstm_input, (h,c) in sequences: | |
| if seq[-1] == end_index: | |
| completed_sequences.append((seq, score)) | |
| continue | |
| lstm_out, (h_new, c_new) = model.lstm.lstm(lstm_input, (h, c)) # lstm_out: (1, 1, 1024) | |
| output = model.lstm.fc(lstm_out.squeeze(1)) # Shape: (1, vocab_size) | |
| log_probs = F.log_softmax(output, dim=1) # Shape: (1, vocab_size) | |
| top_log_probs, top_indices = log_probs.topk(beam_width, dim=1) # Each of shape: (1, beam_width) | |
| for i in range(beam_width): | |
| token = top_indices[0, i].item() | |
| token_log_prob = top_log_probs[0, i].item() | |
| new_seq = seq + [token] | |
| new_score = score + token_log_prob | |
| token_tensor = torch.tensor([token], device=device) | |
| embeddings = self.lstm.embedding(token_tensor) | |
| context, _ = self.lstm.attention(cnn_features, h_new[-1]) | |
| new_lstm_input = torch.cat([embeddings, context], dim=1).unsqueeze(1) | |
| if h_new is not None and c_new is not None: | |
| h_new, c_new = (h_new.clone(), c_new.clone()) | |
| else: | |
| h_new, c_new = None, None | |
| all_candidates.append((new_seq, new_score, new_lstm_input, (h_new, c_new) )) | |
| if not all_candidates: | |
| break | |
| ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True) | |
| sequences = ordered[:beam_width] | |
| if len(completed_sequences) >= beam_width: | |
| break | |
| if len(completed_sequences) == 0: | |
| completed_sequences = sequences | |
| best_seq = max(completed_sequences, key=lambda x: x[1]) | |
| best_caption = decoder(best_seq[0], vocab) | |
| return best_caption | |