Spaces:
Sleeping
Sleeping
| import torch | |
| import string | |
| import torch.nn.functional as F | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| def decoder(indices, vocab): | |
| 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): | |
| """ | |
| Generates captions for images using beam search. | |
| Args: | |
| model (ImgCap): The image captioning model. | |
| images (torch.Tensor): Batch of images. | |
| vocab (Vocab): Vocabulary object. | |
| decoder (function): Function to decode indices to words. | |
| device (str): Device to perform computation on. | |
| start_token (str): Start-of-sequence token. | |
| end_token (str): End-of-sequence token. | |
| beam_width (int): Number of beams to keep. | |
| max_seq_length (int): Maximum length of the generated caption. | |
| Returns: | |
| list: Generated captions for each image in the batch. | |
| """ | |
| 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_feature = model.cnn(images) # Shape: (1, 1024) | |
| lstm_input = model.lstm.projection(cnn_feature).unsqueeze(1) # Shape: (1, 1, 1024) | |
| state = None # Initial LSTM state | |
| # Initialize the beam with the start token | |
| sequences = [([start_index], 0.0, lstm_input, state)] # List of tuples: (sequence, score, input, state) | |
| completed_sequences = [] | |
| for _ in range(max_seq_length): | |
| all_candidates = [] | |
| # Iterate over all current sequences in the beam | |
| for seq, score, lstm_input, state in sequences: | |
| # If the last token is the end token, add the sequence to completed_sequences | |
| if seq[-1] == end_index: | |
| completed_sequences.append((seq, score)) | |
| continue | |
| # Pass the current input and state through the LSTM | |
| lstm_out, state_new = model.lstm.lstm(lstm_input, state) # lstm_out: (1, 1, 1024) | |
| # Pass the LSTM output through the fully connected layer to get logits | |
| output = model.lstm.fc(lstm_out.squeeze(1)) # Shape: (1, vocab_size) | |
| # Compute log probabilities | |
| log_probs = F.log_softmax(output, dim=1) # Shape: (1, vocab_size) | |
| # Get the top beam_width tokens and their log probabilities | |
| top_log_probs, top_indices = log_probs.topk(beam_width, dim=1) # Each of shape: (1, beam_width) | |
| # Iterate over the top tokens to create new candidate sequences | |
| for i in range(beam_width): | |
| token = top_indices[0, i].item() | |
| token_log_prob = top_log_probs[0, i].item() | |
| # Create a new sequence by appending the current token | |
| new_seq = seq + [token] | |
| new_score = score + token_log_prob | |
| # Get the embedding of the new token | |
| token_tensor = torch.tensor([token], device=device) | |
| new_lstm_input = model.lstm.embedding(token_tensor).unsqueeze(1) # Shape: (1, 1, 1024) | |
| # Clone the new state to ensure each beam has its own state | |
| if state_new is not None: | |
| new_state = (state_new[0].clone(), state_new[1].clone()) | |
| else: | |
| new_state = None | |
| # Add the new candidate to all_candidates | |
| all_candidates.append((new_seq, new_score, new_lstm_input, new_state)) | |
| # If no candidates are left to process, break out of the loop | |
| if not all_candidates: | |
| break | |
| # Sort all candidates by score in descending order | |
| ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True) | |
| # Select the top beam_width sequences to form the new beam | |
| sequences = ordered[:beam_width] | |
| # If enough completed sequences are found, stop early | |
| if len(completed_sequences) >= beam_width: | |
| break | |
| # If no sequences have completed, use the current sequences | |
| if len(completed_sequences) == 0: | |
| completed_sequences = sequences | |
| # Select the sequence with the highest score | |
| best_seq, best_score = max(completed_sequences, key=lambda x: x[1]) | |
| if best_seq[0] == start_index: | |
| best_seq = best_seq[1:] | |
| best_caption = decoder(best_seq, vocab) | |
| return best_caption | |
| def generate_caption(model, images, vocab, decoder, device="cpu", start_token="<sos>", end_token="<eos>", max_seq_length=100, top_k=2): | |
| 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) | |
| end_token_appear = {i: False for i in range(batch_size)} | |
| captions = [[] for _ in range(batch_size)] | |
| cnn_feature = model.cnn(images) | |
| lstm_input = model.lstm.projection(cnn_feature).unsqueeze(1) # (B, 1, hidden_size) | |
| state = None | |
| for i in range(max_seq_length): | |
| lstm_out, state = model.lstm.lstm(lstm_input, state) | |
| output = model.lstm.fc(lstm_out.squeeze(1)) | |
| top_k_probs, top_k_indices = torch.topk(F.softmax(output, dim=1), top_k, dim=1) | |
| top_k_probs = top_k_probs / torch.sum(top_k_probs, dim=1, keepdim=True) | |
| top_k_samples = torch.multinomial(top_k_probs, 1).squeeze() | |
| predicted_word_indices = top_k_indices[range(batch_size), top_k_samples] | |
| lstm_input = model.lstm.embedding(predicted_word_indices).unsqueeze(1) # (B, 1, hidden_size) | |
| for j in range(batch_size): | |
| if end_token_appear[j]: | |
| continue | |
| word = vocab.lookup_token(predicted_word_indices[j].item()) | |
| if word == end_token: | |
| end_token_appear[j] = True | |
| captions[j].append(predicted_word_indices[j].item()) | |
| captions = [decoder(caption, vocab) for caption in captions] | |
| return captions | |
| class ResNet50(nn.Module): | |
| def __init__(self): | |
| super(ResNet50, self).__init__() | |
| self.ResNet50 = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) | |
| self.ResNet50.fc = nn.Sequential( | |
| nn.Linear(2048, 1024), | |
| nn.ReLU(), | |
| nn.Dropout(0.5), | |
| nn.Linear(1024, 1024), | |
| nn.ReLU(), | |
| ) | |
| for k,v in self.ResNet50.named_parameters(recurse=True): | |
| if 'fc' in k: | |
| v.requires_grad = True | |
| else: | |
| v.requires_grad = False | |
| def forward(self,x): | |
| return self.ResNet50(x) | |
| ## lSTM (Decoder) | |
| class lstm(nn.Module): | |
| def __init__(self, input_size, hidden_size, number_layers, embedding_dim, vocab_size): | |
| super(lstm, self).__init__() | |
| self.input_size = input_size | |
| self.hidden_size = hidden_size | |
| self.number_layers = number_layers | |
| self.embedding_dim = embedding_dim | |
| self.vocab_size = vocab_size | |
| self.embedding = nn.Embedding(vocab_size, hidden_size) | |
| self.projection = nn.Linear(input_size, hidden_size) | |
| self.relu = nn.ReLU() | |
| self.lstm = nn.LSTM( | |
| input_size=hidden_size, | |
| 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, x, captions): | |
| projected_image = self.projection(x).unsqueeze(dim=1) | |
| embeddings = self.embedding(captions[:, :-1]) | |
| # Concatenate the image feature as frist step with word embeddings | |
| lstm_input = torch.cat((projected_image, embeddings), dim=1) | |
| # print(torch.all(projected_image[:, 0, :] == lstm_input[:, 0, :])) # check | |
| lstm_out, _ = self.lstm(lstm_input) | |
| logits = self.fc(lstm_out) | |
| return logits | |
| ## ImgCap | |
| class ImgCap(nn.Module): | |
| def __init__(self, cnn_feature_size, lstm_hidden_size, num_layers, vocab_size, embedding_dim): | |
| super(ImgCap, self).__init__() | |
| self.cnn = ResNet50() | |
| self.lstm = lstm(input_size=cnn_feature_size, | |
| hidden_size=lstm_hidden_size, | |
| number_layers=num_layers, | |
| embedding_dim=embedding_dim, | |
| vocab_size=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) | |
| end_token_appear = {i: False for i in range(batch_size)} | |
| captions = [[] for _ in range(batch_size)] | |
| cnn_feature = self.cnn(images) | |
| lstm_input = self.lstm.projection(cnn_feature).unsqueeze(1) # (B, 1, hidden_size) | |
| state = None | |
| for i in range(max_seq_length): | |
| lstm_out, state = self.lstm.lstm(lstm_input, state) | |
| output = self.lstm.fc(lstm_out.squeeze(1)) | |
| predicted_word_indices = torch.argmax(output, dim=1) | |
| lstm_input = self.lstm.embedding(predicted_word_indices).unsqueeze(1) # (B, 1, hidden_size) | |
| for j in range(batch_size): | |
| if end_token_appear[j]: | |
| continue | |
| word = vocab.lookup_token(predicted_word_indices[j].item()) | |
| if word == end_token: | |
| end_token_appear[j] = True | |
| captions[j].append(predicted_word_indices[j].item()) | |
| captions = [decoder(caption) for caption in captions] | |
| return captions | |