Duplicate from Vageesh1/clip_gpt2

Co-authored-by: vageesh <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

COCO_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35200360d19ea02ce5c8f007c8bf6d8297e3c16ae3b3fb4b6eeb24ec1c07f8e6
+size 636283447

README.md

@@ -0,0 +1,13 @@
+---
+title: Clip Gpt2
+emoji: 🐨
+colorFrom: pink
+colorTo: indigo
+sdk: streamlit
+sdk_version: 1.19.0
+app_file: app.py
+pinned: false
+duplicated_from: Vageesh1/clip_gpt2
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,84 @@
+import torch 
+import clip
+import PIL.Image
+from PIL import Image
+import skimage.io as io
+import streamlit as st
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
+from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel
+from model import generate2,ClipCaptionModel
+from engine import inference
+
+
+model_trained = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+model_trained.load_state_dict(torch.load('model_trained.pth',map_location=torch.device('cpu')),strict=False)
+image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer       = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
+def show_n_generate(img, model, greedy = True):
+    image = Image.open(img)
+    pixel_values   = image_processor(image, return_tensors ="pt").pixel_values
+
+    if greedy:
+        generated_ids  = model.generate(pixel_values, max_new_tokens = 30)
+    else:
+        generated_ids  = model.generate(
+            pixel_values,
+            do_sample=True,
+            max_new_tokens = 30,
+            top_k=5)
+    generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
+device =  "cpu"
+clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+prefix_length = 10
+
+model = ClipCaptionModel(prefix_length)
+
+model.load_state_dict(torch.load('model.h5',map_location=torch.device('cpu')),strict=False) 
+
+model = model.eval() 
+
+coco_model = ClipCaptionModel(prefix_length)
+coco_model.load_state_dict(torch.load('COCO_model.h5',map_location=torch.device('cpu')),strict=False)
+model = model.eval()  
+
+
+def ui():
+    st.markdown("# Image Captioning")
+    # st.markdown("## Done By- Vageesh and Rushil")
+    uploaded_file = st.file_uploader("Upload an Image", type=['png', 'jpeg', 'jpg'])
+
+    if uploaded_file is not None:
+        image = io.imread(uploaded_file)
+        pil_image = PIL.Image.fromarray(image)
+        image = preprocess(pil_image).unsqueeze(0).to(device)
+
+        option = st.selectbox('Please select the Model',('Clip Captioning','Attention Decoder','VIT+GPT2'))
+
+        if option=='Clip Captioning':
+            with torch.no_grad():
+                prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
+                prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
+            generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed)
+
+            st.image(uploaded_file, width = 500, channels = 'RGB')
+            st.markdown("**PREDICTION:** " + generated_text_prefix)
+        elif option=='Attention Decoder': 
+            out = inference(uploaded_file)
+            st.image(uploaded_file, width = 500, channels = 'RGB')
+            st.markdown("**PREDICTION:** " + out)
+
+        # elif option=='VIT+GPT2': 
+        #     out=show_n_generate(uploaded_file, greedy = False, model = model_trained)
+        #     st.image(uploaded_file, width = 500, channels = 'RGB')
+        #     st.markdown("**PREDICTION:** " + out)
+            
+
+
+if __name__ == '__main__':
+    ui()
+

engine.py

@@ -0,0 +1,42 @@
+import os
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import json
+from neuralnet.model import SeqToSeq
+import wget
+
+url = "https://github.com/Koushik0901/Image-Captioning/releases/download/v1.0/flickr30k.pt"
+# os.system("curl -L https://github.com/Koushik0901/Image-Captioning/releases/download/v1.0/flickr30k.pt")
+filename = wget.download(url)
+
+def inference(img_path):
+    transform = transforms.Compose(
+        [
+            transforms.Resize((299, 299)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ]
+    )
+
+    vocabulary = json.load(open('./vocab.json'))
+
+    model_params = {"embed_size":256, "hidden_size":512, "vocab_size": 7666, "num_layers": 3, "device":"cpu"}
+    model = SeqToSeq(**model_params)
+    checkpoint = torch.load('./flickr30k.pt', map_location = 'cpu')
+    model.load_state_dict(checkpoint['state_dict'])
+
+    img = transform(Image.open(img_path).convert("RGB")).unsqueeze(0)
+
+    result_caption = []
+    model.eval()
+
+    x = model.encoder(img).unsqueeze(0)
+    states = None
+
+    out_captions = model.caption_image(img, vocabulary['itos'], 50)
+    return " ".join(out_captions[1:-1])
+
+
+if __name__ == '__main__':
+    print(inference('./test_examples/dog.png'))

model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a36a09076b9779de2807d3aa533d455a398d70c1250aeb24a5cc9110e3d59a4
+size 636272061

model.py

@@ -0,0 +1,220 @@
+import clip
+import os
+from torch import nn
+import numpy as np
+import torch
+import torch.nn.functional as nnf
+import sys
+from typing import Tuple, List, Union, Optional
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
+from tqdm import tqdm, trange
+import skimage.io as io
+import PIL.Image 
+
+
+N = type(None)
+V = np.array
+ARRAY = np.ndarray
+ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
+VS = Union[Tuple[V, ...], List[V]]
+VN = Union[V, N]
+VNS = Union[VS, N]
+T = torch.Tensor
+TS = Union[Tuple[T, ...], List[T]]
+TN = Optional[T]
+TNS = Union[Tuple[TN, ...], List[TN]]
+TSN = Optional[TS]
+TA = Union[T, ARRAY]
+
+
+D = torch.device
+
+def get_device(device_id: int) -> D:
+    if not torch.cuda.is_available():
+        return CPU
+    device_id = min(torch.cuda.device_count() - 1, device_id)
+    return torch.device(f'cuda:{device_id}')
+
+
+CUDA = get_device
+
+current_directory = os.getcwd()
+save_path = os.path.join(os.path.dirname(current_directory), "pretrained_models")
+os.makedirs(save_path, exist_ok=True)
+model_path = os.path.join(save_path, 'model_wieghts.pt')
+
+
+class MLP(nn.Module):
+
+    def forward(self, x: T) -> T:
+        return self.model(x)
+
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+        layers = []
+        for i in range(len(sizes) -1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+class ClipCaptionModel(nn.Module):
+
+    #@functools.lru_cache #FIXME
+    def get_dummy_token(self, batch_size: int, device: D) -> T:
+        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
+
+    def forward(self, tokens: T, prefix: T, mask: Optional[T] = None, labels: Optional[T] = None):
+        embedding_text = self.gpt.transformer.wte(tokens)
+        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
+        #print(embedding_text.size()) #torch.Size([5, 67, 768])
+        #print(prefix_projections.size()) #torch.Size([5, 1, 768])
+        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
+        if labels is not None:
+            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
+            labels = torch.cat((dummy_token, tokens), dim=1)
+        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
+        return out
+
+    def __init__(self, prefix_length: int, prefix_size: int = 512):
+        super(ClipCaptionModel, self).__init__()
+        self.prefix_length = prefix_length
+        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+        if prefix_length > 10:  # not enough memory
+            self.clip_project = nn.Linear(prefix_size, self.gpt_embedding_size * prefix_length)
+        else:
+            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2, self.gpt_embedding_size * prefix_length))
+
+
+class ClipCaptionPrefix(ClipCaptionModel):
+
+    def parameters(self, recurse: bool = True):
+        return self.clip_project.parameters()
+
+    def train(self, mode: bool = True):
+        super(ClipCaptionPrefix, self).train(mode)
+        self.gpt.eval()
+        return self
+    
+def generate_beam(model, tokenizer, beam_size: int = 5, prompt=None, embed=None,
+                  entry_length=67, temperature=1., stop_token: str = '.'):
+
+    model.eval()
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    tokens = None
+    scores = None
+    device = next(model.parameters()).device
+    seq_lengths = torch.ones(beam_size, device=device)
+    is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)
+    with torch.no_grad():
+        if embed is not None:
+            generated = embed
+        else:
+            if tokens is None:
+                tokens = torch.tensor(tokenizer.encode(prompt))
+                tokens = tokens.unsqueeze(0).to(device)
+                generated = model.gpt.transformer.wte(tokens)
+        for i in range(entry_length):
+            outputs = model.gpt(inputs_embeds=generated)
+            logits = outputs.logits
+            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+            logits = logits.softmax(-1).log()
+            if scores is None:
+                scores, next_tokens = logits.topk(beam_size, -1)
+                generated = generated.expand(beam_size, *generated.shape[1:])
+                next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)
+                if tokens is None:
+                    tokens = next_tokens
+                else:
+                    tokens = tokens.expand(beam_size, *tokens.shape[1:])
+                    tokens = torch.cat((tokens, next_tokens), dim=1)
+            else:
+                logits[is_stopped] = -float(np.inf)
+                logits[is_stopped, 0] = 0
+                scores_sum = scores[:, None] + logits
+                seq_lengths[~is_stopped] += 1
+                scores_sum_average = scores_sum / seq_lengths[:, None]
+                scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(beam_size, -1)
+                next_tokens_source = next_tokens // scores_sum.shape[1]
+                seq_lengths = seq_lengths[next_tokens_source]
+                next_tokens = next_tokens % scores_sum.shape[1]
+                next_tokens = next_tokens.unsqueeze(1)
+                tokens = tokens[next_tokens_source]
+                tokens = torch.cat((tokens, next_tokens), dim=1)
+                generated = generated[next_tokens_source]
+                scores = scores_sum_average * seq_lengths
+                is_stopped = is_stopped[next_tokens_source]
+            next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(generated.shape[0], 1, -1)
+            generated = torch.cat((generated, next_token_embed), dim=1)
+            is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()
+            if is_stopped.all():
+                break
+    scores = scores / seq_lengths
+    output_list = tokens.cpu().numpy()
+    output_texts = [tokenizer.decode(output[:int(length)]) for output, length in zip(output_list, seq_lengths)]
+    order = scores.argsort(descending=True)
+    output_texts = [output_texts[i] for i in order]
+    return output_texts
+
+def generate2(
+        model,
+        tokenizer,
+        tokens=None,
+        prompt=None,
+        embed=None,
+        entry_count=1,
+        entry_length=67,  # maximum number of words
+        top_p=0.8,
+        temperature=1.,
+        stop_token: str = '.',
+):
+    model.eval()
+    generated_num = 0
+    generated_list = []
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    filter_value = -float("Inf")
+    device = next(model.parameters()).device
+
+    with torch.no_grad():
+
+        for entry_idx in trange(entry_count):
+            if embed is not None:
+                generated = embed
+            else:
+                if tokens is None:
+                    tokens = torch.tensor(tokenizer.encode(prompt))
+                    tokens = tokens.unsqueeze(0).to(device)
+
+                generated = model.gpt.transformer.wte(tokens)
+
+            for i in range(entry_length):
+
+                outputs = model.gpt(inputs_embeds=generated)
+                logits = outputs.logits
+                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(nnf.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                                                    ..., :-1
+                                                    ].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                logits[:, indices_to_remove] = filter_value
+                next_token = torch.argmax(logits, -1).unsqueeze(0)
+                next_token_embed = model.gpt.transformer.wte(next_token)
+                if tokens is None:
+                    tokens = next_token
+                else:
+                    tokens = torch.cat((tokens, next_token), dim=1)
+                generated = torch.cat((generated, next_token_embed), dim=1)
+                if stop_token_index == next_token.item():
+                    break
+
+            output_list = list(tokens.squeeze().cpu().numpy())
+            output_text = tokenizer.decode(output_list)
+            generated_list.append(output_text)
+
+    return generated_list[0]

model_trained.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f44c397a407f1687578a0346cbe19262b4ba6954c3256ec656ade873ac57d07
+size 982140285

neuralnet/dataset.py

@@ -0,0 +1,139 @@
+import os  # when loading file paths
+import pandas as pd  # for lookup in annotation file
+import spacy  # for tokenizer
+import torch
+from torch.nn.utils.rnn import pad_sequence  # pad batch
+from torch.utils.data import DataLoader, Dataset
+from PIL import Image  # Load img
+import torchvision.transforms as transforms
+import json
+
+# 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.stoi)
+
+    @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_name"]
+        self.captions = self.df["comment"]
+
+        # 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)
+        targets = [item[1] for item in batch]
+        targets = pad_sequence(targets, batch_first=False, padding_value=self.pad_idx)
+
+        return imgs, targets
+
+
+def get_loader(
+    root_folder,
+    annotation_file,
+    transform,
+    batch_size=64,
+    num_workers=2,
+    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(
+        "/home/koushik/vscode/Projects/pytorch/img2text_v1/flickr30k/flickr30k_images/", "/home/koushik/vscode/Projects/pytorch/img2text_v1/flickr30k/results.csv", transform=transform
+    )
+
+    for idx, (imgs, captions) in enumerate(loader):
+        print(imgs.shape)
+        print(captions.shape)
+        print(len(dataset.vocab))
+        test = {"itos":dataset.vocab.itos, "stoi": dataset.vocab.stoi}
+        json.dump(test, open('test.json', 'w'))
+        break

neuralnet/model.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+
+class InceptionEncoder(nn.Module):
+    def __init__(self, embed_size, train_CNN=False):
+        super(InceptionEncoder, self).__init__()
+        self.train_CNN = train_CNN
+        self.inception = models.inception_v3(pretrained=True, aux_logits=False)
+        self.inception.fc = nn.Linear(self.inception.fc.in_features, embed_size)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm1d(embed_size, momentum = 0.01)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, images):
+        features = self.inception(images)
+        norm_features = self.bn(features)
+        return self.dropout(self.relu(norm_features))
+
+
+class LstmDecoder(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, device = 'cpu'):
+        super(LstmDecoder, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+        self.device = device
+        self.embed = nn.Embedding(vocab_size, embed_size)
+        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers = self.num_layers)
+        self.linear = nn.Linear(hidden_size, vocab_size)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, encoder_out, captions):
+        h0 = torch.zeros(self.num_layers, encoder_out.shape[0], self.hidden_size).to(self.device).requires_grad_()
+        c0 = torch.zeros(self.num_layers, encoder_out.shape[0], self.hidden_size).to(self.device).requires_grad_()
+        embeddings = self.dropout(self.embed(captions))
+        embeddings = torch.cat((encoder_out.unsqueeze(0), embeddings), dim=0)
+        hiddens, (hn, cn) = self.lstm(embeddings, (h0.detach(), c0.detach()))
+        outputs = self.linear(hiddens)
+        return outputs
+
+
+class SeqToSeq(nn.Module):
+    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, device = 'cpu'):
+        super(SeqToSeq, self).__init__()
+        self.encoder = InceptionEncoder(embed_size)
+        self.decoder = LstmDecoder(embed_size, hidden_size, vocab_size, num_layers, device)
+
+    def forward(self, images, captions):
+        features = self.encoder(images)
+        outputs = self.decoder(features, captions)
+        return outputs
+
+    def caption_image(self, image, vocabulary, max_length = 50):
+        result_caption = []
+
+        with torch.no_grad():
+            x = self.encoder(image).unsqueeze(0)
+            states = None
+
+            for _ in range(max_length):
+                hiddens, states = self.decoder.lstm(x, states)
+                output = self.decoder.linear(hiddens.squeeze(0))
+                predicted = output.argmax(1)
+                result_caption.append(predicted.item())
+                x = self.decoder.embed(predicted).unsqueeze(0)
+
+                if vocabulary[str(predicted.item())] == "<EOS>":
+                    break
+
+        return [vocabulary[str(idx)] for idx in result_caption]

neuralnet/train.py

@@ -0,0 +1,130 @@
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+from torch.utils.tensorboard import SummaryWriter # For TensorBoard
+from utils import save_checkpoint, load_checkpoint, print_examples
+from dataset import get_loader
+from model import SeqToSeq
+from tabulate import tabulate # To tabulate loss and epoch
+import argparse
+import json
+
+def main(args):
+    transform = transforms.Compose(
+        [
+            transforms.Resize((356, 356)),
+            transforms.RandomCrop((299, 299)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ]
+    )
+
+    train_loader, _ = get_loader(
+        root_folder = args.root_dir,
+        annotation_file = args.csv_file,
+        transform=transform,
+        batch_size = 64,
+        num_workers=2,
+    )
+    vocab = json.load(open('vocab.json'))
+
+    torch.backends.cudnn.benchmark = True
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    load_model = False
+    save_model = True
+    train_CNN = False
+
+    # Hyperparameters
+    embed_size = args.embed_size
+    hidden_size = args.hidden_size
+    vocab_size = len(vocab['stoi'])
+    num_layers = args.num_layers
+    learning_rate = args.lr
+    num_epochs = args.num_epochs
+    # for tensorboard
+
+    
+    writer = SummaryWriter(args.log_dir)
+    step = 0
+    model_params = {'embed_size': embed_size, 'hidden_size': hidden_size, 'vocab_size':vocab_size, 'num_layers':num_layers}
+    # initialize model, loss etc
+    model = SeqToSeq(**model_params, device = device).to(device)
+    criterion = nn.CrossEntropyLoss(ignore_index = vocab['stoi']["<PAD>"])
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+    # Only finetune the CNN
+    for name, param in model.encoder.inception.named_parameters():
+        if "fc.weight" in name or "fc.bias" in name:
+            param.requires_grad = True
+        else:
+            param.requires_grad = train_CNN
+
+    #load from a save checkpoint
+    if load_model:
+        step = load_checkpoint(torch.load(args.save_path), model, optimizer)
+
+    model.train()
+    best_loss, best_epoch = 10, 0
+    for epoch in range(num_epochs):
+        print_examples(model, device, vocab['itos'])
+
+        for idx, (imgs, captions) in tqdm(
+            enumerate(train_loader), total=len(train_loader), leave=False):
+            imgs = imgs.to(device)
+            captions = captions.to(device)
+
+            outputs = model(imgs, captions[:-1])
+            loss = criterion(
+                outputs.reshape(-1, outputs.shape[2]), captions.reshape(-1)
+            )
+
+            writer.add_scalar("Training loss", loss.item(), global_step=step)
+            step += 1
+
+            optimizer.zero_grad()
+            loss.backward(loss)
+            optimizer.step()
+		
+        train_loss = loss.item()
+        if train_loss < best_loss:
+            best_loss = train_loss
+            best_epoch = epoch + 1
+            if save_model:
+                checkpoint = {
+                    "model_params": model_params,
+		            "state_dict": model.state_dict(),
+		            "optimizer": optimizer.state_dict(),
+		            "step": step
+		        }
+                save_checkpoint(checkpoint, args.save_path)
+
+
+        table = [["Loss:", train_loss],
+				["Step:", step],
+                ["Epoch:", epoch + 1],
+		 		["Best Loss:", best_loss],
+		  		["Best Epoch:", best_epoch]]
+        print(tabulate(table))
+	
+	
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--root_dir', type = str, default = './flickr30k/flickr30k_images', help = 'path to images folder')
+    parser.add_argument('--csv_file', type = str, default = './flickr30k/results.csv', help = 'path to captions csv file')
+    parser.add_argument('--log_dir', type = str, default = './drive/MyDrive/TensorBoard/', help = 'path to save tensorboard logs')
+    parser.add_argument('--save_path', type = str, default = './drive/MyDrive/checkpoints/Seq2Seq.pt', help = 'path to save checkpoint')
+    # Model Params
+    parser.add_argument('--batch_size', type = int, default = 64)
+    parser.add_argument('--num_epochs', type = int, default = 100)
+    parser.add_argument('--embed_size', type = int, default=256)
+    parser.add_argument('--hidden_size', type = int, default=512)
+    parser.add_argument('--lr', type = float, default= 0.001)
+    parser.add_argument('--num_layers', type = int, default = 3, help = 'number of lstm layers')
+
+    args = parser.parse_args()
+    
+    main(args)

neuralnet/utils.py

@@ -0,0 +1,42 @@
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+
+
+def print_examples(model, device, vocab):
+    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.png").convert("RGB")).unsqueeze(0)
+    print("dog.png PREDICTION: " + " ".join(model.caption_image(test_img1.to(device), vocab)))
+
+    test_img2 = transform(Image.open("./test_examples/dirt_bike.png").convert("RGB")).unsqueeze(0)
+    print("dirt_bike.png PREDICTION: " + " ".join(model.caption_image(test_img2.to(device), vocab)))
+
+    test_img3 = transform(Image.open("./test_examples/surfing.png").convert("RGB")).unsqueeze(0)
+    print("wave.png PREDICTION: " + " ".join(model.caption_image(test_img3.to(device), vocab)))
+
+    test_img4 = transform(Image.open("./test_examples/horse.png").convert("RGB")).unsqueeze(0)
+    print("horse.png PREDICTION: " + " ".join(model.caption_image(test_img4.to(device), vocab)))
+    
+    test_img5 = transform(Image.open("./test_examples/camera.png").convert("RGB")).unsqueeze(0)
+    print("camera.png PREDICTION: " + " ".join(model.caption_image(test_img5.to(device), vocab)))
+    model.train()
+
+
+def save_checkpoint(state, filename="/content/drive/MyDrive/checkpoints/Seq2Seq.pt"):
+    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

requirements.txt

@@ -0,0 +1,19 @@
+torch
+torchvision
+ftfy 
+git+https://github.com/openai/CLIP.git
+regex 
+tqdm
+streamlit
+scikit-image
+pillow
+pandas
+transformers
+numpy
+spacy
+tqdm
+tabulate
+click==7.1.1
+gdown
+wget
+altair<5