leonhardt/ImageCaptionDemo

1import os  # when loading file paths2 3import pandas as pd  # for lookup in annotation file4import spacy  # for tokenizer5import torch6import torchvision.transforms as transforms7from PIL import Image  # Load img8from torch.nn.utils.rnn import pad_sequence  # pad batch9from torch.utils.data import DataLoader, Dataset10 11# We want to convert text -> numerical values12# 1. We need a Vocabulary mapping each word to a index13# 2. We need to setup a Pytorch dataset to load the data14# 3. Setup padding of every batch (all examples should be15#    of same seq_len and setup dataloader)16 17# Download with: python -m spacy download en18spacy_eng = spacy.load("en_core_web_sm")19 20 21class Vocabulary:22    def __init__(self, freq_threshold):23        self.itos = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}24        self.stoi = {"<PAD>": 0, "<SOS>": 1, "<EOS>": 2, "<UNK>": 3}25        self.freq_threshold = freq_threshold26 27    def __len__(self):28        return len(self.itos)29 30    @staticmethod31    def tokenizer_eng(text):32        return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]33 34    def build_vocabulary(self, sentence_list):35        frequencies = {}36        idx = 437 38        for sentence in sentence_list:39            for word in self.tokenizer_eng(sentence):40                if word not in frequencies:41                    frequencies[word] = 142 43                else:44                    frequencies[word] += 145 46                if frequencies[word] == self.freq_threshold:47                    self.stoi[word] = idx48                    self.itos[idx] = word49                    idx += 150 51    def numericalize(self, text):52        tokenized_text = self.tokenizer_eng(text)53 54        return [55            self.stoi[token] if token in self.stoi else self.stoi["<UNK>"]56            for token in tokenized_text57        ]58 59 60class FlickrDataset(Dataset):61    def __init__(self, root_dir, captions_file, transform=None, freq_threshold=5):62        self.root_dir = root_dir63        self.df = pd.read_csv(captions_file)64        self.transform = transform65 66        # Get img, caption columns67        self.imgs = self.df["image"]68        self.captions = self.df["caption"]69 70        # Initialize vocabulary and build vocab71        self.vocab = Vocabulary(freq_threshold)72        self.vocab.build_vocabulary(self.captions.tolist())73 74    def __len__(self):75        return len(self.df)76 77    def __getitem__(self, index):78        caption = self.captions[index]79        img_id = self.imgs[index]80        img = Image.open(os.path.join(self.root_dir, img_id)).convert("RGB")81 82        if self.transform is not None:83            img = self.transform(img)84 85        numericalized_caption = [self.vocab.stoi["<SOS>"]]86        numericalized_caption += self.vocab.numericalize(caption)87        numericalized_caption.append(self.vocab.stoi["<EOS>"])88 89        return img, torch.tensor(numericalized_caption)90 91 92class MyCollate:93    def __init__(self, pad_idx):94        self.pad_idx = pad_idx95 96    def __call__(self, batch):97        imgs = [item[0].unsqueeze(0) for item in batch]98        imgs = torch.cat(imgs, dim=0) # [BCHW]99        targets = [item[1] for item in batch]  # [BL] L长度不同100        targets = pad_sequence(targets, batch_first=False,  # [LB] L长度相同101                               padding_value=self.pad_idx)102        return imgs, targets103 104 105def get_loader(106    root_folder,107    annotation_file,108    transform,109    batch_size=32,110    num_workers=8,111    shuffle=True,112    pin_memory=True,113):114    dataset = FlickrDataset(root_folder, annotation_file, transform=transform)115 116    pad_idx = dataset.vocab.stoi["<PAD>"]117 118    loader = DataLoader(119        dataset=dataset,120        batch_size=batch_size,121        num_workers=num_workers,122        shuffle=shuffle,123        pin_memory=pin_memory,124        collate_fn=MyCollate(pad_idx=pad_idx),125    )126 127    return loader, dataset128 129 130if __name__ == "__main__":131    transform = transforms.Compose(132        [transforms.Resize((224, 224)), transforms.ToTensor(),]133    )134 135    loader, dataset = get_loader(136        "flickr8k/images/", "flickr8k/captions.txt", transform=transform137    )138 139    for idx, (imgs, captions) in enumerate(loader):140        print(imgs.shape)141        print(captions.shape)142