add app

app.py

@@ -0,0 +1,51 @@
+import gradio as gr
+import pandas as pd
+import torch
+from torchvision.transforms import transforms
+
+from get_loader import Vocabulary
+from model import CNNtoRNN
+
+
+def predict(img):
+    img = transform(img)
+    output = model.caption_image(img.unsqueeze(0).to(device), vocab)
+    return " ".join(output[1:-1])
+
+if __name__ == '__main__':
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = CNNtoRNN(
+        embed_size=256,
+        hidden_size=256,
+        vocab_size=2994,
+        num_layers=1
+    )
+    print('Loading model weights...')
+    model.load_state_dict(torch.load(
+        'my_checkpoint.pth.tar', map_location=device)["state_dict"])
+    model.to(device)
+    model.eval()
+    print('Building vocabulary...')
+    vocab = Vocabulary(5)
+    df = pd.read_csv('captions.txt')
+    vocab.build_vocabulary(df['caption'].tolist())
+
+    transform = transforms.Compose(
+        [
+            transforms.ToPILImage(),
+            transforms.Resize((299, 299)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]
+    )
+    print('Creating app...')
+    app = gr.Interface(
+        fn=predict,
+        inputs=gr.Image(shape=(256, 256)),
+        outputs="text",
+    )
+    print('done!')
+
+    app.launch(
+        share=True,
+    )

get_loader.py

@@ -0,0 +1,141 @@
+import os  # when loading file paths
+
+import pandas as pd  # for lookup in annotation file
+import spacy  # for tokenizer
+import torch
+import torchvision.transforms as transforms
+from PIL import Image  # Load img
+from torch.nn.utils.rnn import pad_sequence  # pad batch
+from torch.utils.data import DataLoader, Dataset
+
+# We want to convert text -> numerical values
+# 1. We need a Vocabulary mapping each word to a index
+# 2. We need to setup a Pytorch dataset to load the data
+# 3. Setup padding of every batch (all examples should be
+#    of same seq_len and setup dataloader)
+
+# Download with: python -m spacy download en
+spacy_eng = spacy.load("en_core_web_sm")
+
+
+class Vocabulary:
+    def __init__(self, freq_threshold):
+        self.itos = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}
+        self.stoi = {"<PAD>": 0, "<SOS>": 1, "<EOS>": 2, "<UNK>": 3}
+        self.freq_threshold = freq_threshold
+
+    def __len__(self):
+        return len(self.itos)
+
+    @staticmethod
+    def tokenizer_eng(text):
+        return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]
+
+    def build_vocabulary(self, sentence_list):
+        frequencies = {}
+        idx = 4
+
+        for sentence in sentence_list:
+            for word in self.tokenizer_eng(sentence):
+                if word not in frequencies:
+                    frequencies[word] = 1
+
+                else:
+                    frequencies[word] += 1
+
+                if frequencies[word] == self.freq_threshold:
+                    self.stoi[word] = idx
+                    self.itos[idx] = word
+                    idx += 1
+
+    def numericalize(self, text):
+        tokenized_text = self.tokenizer_eng(text)
+
+        return [
+            self.stoi[token] if token in self.stoi else self.stoi["<UNK>"]
+            for token in tokenized_text
+        ]
+
+
+class FlickrDataset(Dataset):
+    def __init__(self, root_dir, captions_file, transform=None, freq_threshold=5):
+        self.root_dir = root_dir
+        self.df = pd.read_csv(captions_file)
+        self.transform = transform
+
+        # Get img, caption columns
+        self.imgs = self.df["image"]
+        self.captions = self.df["caption"]
+
+        # Initialize vocabulary and build vocab
+        self.vocab = Vocabulary(freq_threshold)
+        self.vocab.build_vocabulary(self.captions.tolist())
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, index):
+        caption = self.captions[index]
+        img_id = self.imgs[index]
+        img = Image.open(os.path.join(self.root_dir, img_id)).convert("RGB")
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        numericalized_caption = [self.vocab.stoi["<SOS>"]]
+        numericalized_caption += self.vocab.numericalize(caption)
+        numericalized_caption.append(self.vocab.stoi["<EOS>"])
+
+        return img, torch.tensor(numericalized_caption)
+
+
+class MyCollate:
+    def __init__(self, pad_idx):
+        self.pad_idx = pad_idx
+
+    def __call__(self, batch):
+        imgs = [item[0].unsqueeze(0) for item in batch]
+        imgs = torch.cat(imgs, dim=0) # [BCHW]
+        targets = [item[1] for item in batch]  # [BL] L长度不同
+        targets = pad_sequence(targets, batch_first=False,  # [LB] L长度相同
+                               padding_value=self.pad_idx)
+        return imgs, targets
+
+
+def get_loader(
+    root_folder,
+    annotation_file,
+    transform,
+    batch_size=32,
+    num_workers=8,
+    shuffle=True,
+    pin_memory=True,
+):
+    dataset = FlickrDataset(root_folder, annotation_file, transform=transform)
+
+    pad_idx = dataset.vocab.stoi["<PAD>"]
+
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        shuffle=shuffle,
+        pin_memory=pin_memory,
+        collate_fn=MyCollate(pad_idx=pad_idx),
+    )
+
+    return loader, dataset
+
+
+if __name__ == "__main__":
+    transform = transforms.Compose(
+        [transforms.Resize((224, 224)), transforms.ToTensor(),]
+    )
+
+    loader, dataset = get_loader(
+        "flickr8k/images/", "flickr8k/captions.txt", transform=transform
+    )
+
+    for idx, (imgs, captions) in enumerate(loader):
+        print(imgs.shape)
+        print(captions.shape)

model.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+#----------------------------------------------------------------------------
+
+class EncoderCNN(nn.Module):
+    def __init__(self, embed_size) -> None:
+        super().__init__()
+        self.inception = models.inception_v3(pretrained=True, aux_logits=False)
+        for param in self.inception.parameters():
+            param.requires_grad = False
+        self.inception.fc = nn.Linear(self.inception.fc.in_features, embed_size)
+        self.relu = nn.ReLU(True)
+        self.times = []
+        self.dropout = nn.Dropout(0.5)
+    
+    def forward(self, imgs):
+        features = self.inception(imgs)
+        return self.dropout(self.relu(features))
+
+#----------------------------------------------------------------------------
+
+class DecoderRNN(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers) -> None:
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_size)
+        self.LSTM = nn.LSTM(embed_size, hidden_size, num_layers)
+        self.linear = nn.Linear(hidden_size, vocab_size)
+        self.dropout = nn.Dropout(0.5)
+    
+    def forward(self, features, captions):
+        embbedings = self.dropout(self.embed(captions))
+        # unsqueeze(0) 添加时间维度seq_len
+        embbedings = torch.cat([features.unsqueeze(0), embbedings], dim=0)
+        hiddens, _ = self.LSTM(embbedings)
+        outputs = self.linear(hiddens)
+        return outputs
+
+#----------------------------------------------------------------------------
+
+class CNNtoRNN(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers) -> None:
+        super().__init__()
+        self.encoderCNN = EncoderCNN(embed_size)
+        self.decoderRNN = DecoderRNN(embed_size, hidden_size, vocab_size, num_layers)
+    
+    def forward(self, imgs, captions):
+        features = self.encoderCNN(imgs)
+        outputs = self.decoderRNN(features, captions)
+        return outputs
+    
+    def caption_image(self, img, vocab, max_length=50):
+        result_caption = []
+
+        with torch.no_grad():
+            x = self.encoderCNN(img).unsqueeze(0)
+            states = None
+
+            for _ in range(max_length):
+                # 逐个预测
+                h, states = self.decoderRNN.LSTM(x, states)
+                output = self.decoderRNN.linear(h.squeeze(0))
+                predicted = output.argmax(1)
+                # 预测的值作为下一次预测的输入
+                result_caption.append(predicted.item())
+                x = self.decoderRNN.embed(predicted).unsqueeze(0)
+
+                if vocab.itos[predicted.item()] == '<EOS>':
+                    break
+
+        return [vocab.itos[idx] for idx in result_caption]

requirements.txt

@@ -0,0 +1,2 @@
+spacy
+pandas

utils.py

@@ -0,0 +1,69 @@
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+
+
+def print_examples(model, device, dataset):
+    transform = transforms.Compose(
+        [
+            transforms.Resize((299, 299)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]
+    )
+
+    model.eval()
+    test_img1 = transform(Image.open("test_examples/dog.jpg").convert("RGB")).unsqueeze(
+        0
+    )
+    print("Example 1 CORRECT: Dog on a beach by the ocean")
+    print(
+        "Example 1 OUTPUT: "
+        + " ".join(model.caption_image(test_img1.to(device), dataset.vocab))
+    )
+    test_img2 = transform(
+        Image.open("test_examples/child.jpg").convert("RGB")
+    ).unsqueeze(0)
+    print("Example 2 CORRECT: Child holding red frisbee outdoors")
+    print(
+        "Example 2 OUTPUT: "
+        + " ".join(model.caption_image(test_img2.to(device), dataset.vocab))
+    )
+    test_img3 = transform(Image.open("test_examples/bus.png").convert("RGB")).unsqueeze(
+        0
+    )
+    print("Example 3 CORRECT: Bus driving by parked cars")
+    print(
+        "Example 3 OUTPUT: "
+        + " ".join(model.caption_image(test_img3.to(device), dataset.vocab))
+    )
+    test_img4 = transform(
+        Image.open("test_examples/boat.png").convert("RGB")
+    ).unsqueeze(0)
+    print("Example 4 CORRECT: A small boat in the ocean")
+    print(
+        "Example 4 OUTPUT: "
+        + " ".join(model.caption_image(test_img4.to(device), dataset.vocab))
+    )
+    test_img5 = transform(
+        Image.open("test_examples/horse.png").convert("RGB")
+    ).unsqueeze(0)
+    print("Example 5 CORRECT: A cowboy riding a horse in the desert")
+    print(
+        "Example 5 OUTPUT: "
+        + " ".join(model.caption_image(test_img5.to(device), dataset.vocab))
+    )
+    model.train()
+
+
+def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
+    print("=> Saving checkpoint")
+    torch.save(state, filename)
+
+
+def load_checkpoint(checkpoint, model, optimizer):
+    print("=> Loading checkpoint")
+    model.load_state_dict(checkpoint["state_dict"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    step = checkpoint["step"]
+    return step