{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch Resnet to get image features then LSTM with attention to generate text\nFeel free to leave any comments or questions"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n\n        \nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport PIL\nfrom typing import Any, Callable, cast, Dict, List, Optional, Tuple\nfrom torchvision.transforms.transforms import Compose, Normalize, Resize, ToTensor, RandomHorizontalFlip, RandomCrop\nfrom torch.nn.utils.rnn import pad_sequence\nimport torch.optim as optim\nimport torch.nn.functional as F\n\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n\n\nimport os\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\ntqdm.pandas()\n#import Levenshtein\n#import cv2\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\nimport seaborn as sns\nimport time\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#make train and test as in https://www.kaggle.com/yasufuminakama/molecular-translation-naive-baseline\ntrain = pd.read_csv('../input/bms-molecular-translation/train_labels.csv')\ntest = pd.read_csv('../input/bms-molecular-translation/sample_submission.csv')\n\ndef get_train_file_path(image_id):\n    return \"../input/bms-molecular-translation/train/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )\n\ndef get_test_file_path(image_id):\n    return \"../input/bms-molecular-translation/test/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )\n\ntrain['file_path'] = train['image_id'].progress_apply(get_train_file_path)\ntest['file_path'] = test['image_id'].progress_apply(get_test_file_path)\n\n\nprint(f'train.shape: {train.shape}  test.shape: {test.shape}')\ndisplay(train.head())\ndisplay(test.head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#make vocab\n\nwords=set()\nfor st in train['InChI']:\n    words.update(set(st))\nlen(words)\n\nvocab=list(words)\nvocab.append('<sos>')\nvocab.append('<eos>')\nvocab.append('<pad>')\nstoi={'C': 0,')': 1,'P': 2,'l': 3,'=': 4,'3': 5,'N': 6,'I': 7,'2': 8,'6': 9,'H': 10,'4': 11,'F': 12,'0': 13,'1': 14,'-': 15,'O': 16,'8': 17,\n ',': 18,'B': 19,'(': 20,'7': 21,'r': 22,'/': 23,'m': 24,'c': 25,'s': 26,'h': 27,'i': 28,'t': 29,'T': 30,'n': 31,'5': 32,'+': 33,'b': 34,'9': 35,\n 'D': 36,'S': 37,'<sos>': 38,'<eos>': 39,'<pad>': 40}\nitos={item[1]:item[0] for item in stoi.items()}\n\n\ndef string_to_ints(string):\n    l=[stoi['<sos>']]\n    for s in string:\n        l.append(stoi[s])\n    l.append(stoi['<eos>'])\n    return l\ndef ints_to_string(l):\n    return ''.join(list(map(lambda i:itos[i],l)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def pil_loader(path: str) -> Image.Image: #copied from torchvision\n    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n    with open(path, 'rb') as f:\n        img = Image.open(f)\n        return img.convert('RGB')\n    \ndef default_loader(path: str) -> Any:\n    from torchvision import get_image_backend\n    if get_image_backend() == 'accimage':\n        return accimage_loader(path)\n    else:\n        return pil_loader(path)\n    \n\n    \nclass InputDatasetTest(Dataset):\n    def __init__(self,paths,transform):\n        self.paths=paths\n        self.loader=default_loader\n        self.transform=transform\n    def __len__(self):\n        return len(self.paths)\n    def __getitem__(self,idx):\n        sample=self.loader(self.paths[idx])\n        sample=self.transform(sample)\n        return sample,idx\n    \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#model adapted from https://www.kaggle.com/mdteach/image-captioning-with-attention-pytorch/data\nclass Attention(nn.Module):\n    def __init__(self, encoder_dim,decoder_dim,attention_dim):\n        super(Attention, self).__init__()\n        \n        self.attention_dim = attention_dim\n        \n        self.W = nn.Linear(decoder_dim,attention_dim)\n        self.U = nn.Linear(encoder_dim,attention_dim)\n        \n        self.A = nn.Linear(attention_dim,1)\n        \n        \n        \n        \n    def forward(self, features, hidden_state):\n        u_hs = self.U(features)     #(batch_size,64,attention_dim)\n        w_ah = self.W(hidden_state) #(batch_size,attention_dim)\n        \n        combined_states = torch.tanh(u_hs + w_ah.unsqueeze(1)) #(batch_size,64,attemtion_dim)\n        \n        attention_scores = self.A(combined_states)         #(batch_size,64,1)\n        attention_scores = attention_scores.squeeze(2)     #(batch_size,64)\n        \n        \n        alpha = F.softmax(attention_scores,dim=1)          #(batch_size,64)\n        \n        attention_weights = features * alpha.unsqueeze(2)  #(batch_size,64,features_dim)\n        attention_weights = attention_weights.sum(dim=1)   #(batch_size,64)\n        \n        return alpha,attention_weights\n\nclass DecoderRNN(nn.Module):\n    def __init__(self,embed_size, vocab_size, attention_dim,encoder_dim,decoder_dim,drop_prob=0.3):\n        super().__init__()\n        \n        #save the model param\n        self.vocab_size = vocab_size\n        self.attention_dim = attention_dim\n        self.decoder_dim = decoder_dim\n        \n        self.embedding = nn.Embedding(vocab_size,embed_size)\n        self.attention = Attention(encoder_dim,decoder_dim,attention_dim)\n        \n        \n        self.init_h = nn.Linear(encoder_dim, decoder_dim)  \n        self.init_c = nn.Linear(encoder_dim, decoder_dim)  \n        self.lstm_cell = nn.LSTMCell(embed_size+encoder_dim,decoder_dim,bias=True)\n        self.f_beta = nn.Linear(decoder_dim, encoder_dim)\n        \n        \n        self.fcn = nn.Linear(decoder_dim,vocab_size)\n        self.drop = nn.Dropout(drop_prob)\n        \n        \n    \n    def forward(self, features, captions):\n        \n        #vectorize the caption\n        embeds = self.embedding(captions)\n        \n        # Initialize LSTM state\n        h, c = self.init_hidden_state(features)  # (batch_size, decoder_dim)\n        \n        #get the seq length to iterate\n        seq_length = len(captions[0])-1 #Exclude the last one\n        batch_size = captions.size(0)\n        num_features = features.size(1)\n        \n        preds = torch.zeros(batch_size, seq_length, self.vocab_size).to(device)\n        alphas = torch.zeros(batch_size, seq_length,num_features).to(device)\n                \n        for s in range(seq_length):\n            alpha,context = self.attention(features, h)\n            lstm_input = torch.cat((embeds[:, s], context), dim=1)\n            h, c = self.lstm_cell(lstm_input, (h, c))\n                    \n            output = self.fcn(self.drop(h))\n            \n            preds[:,s] = output\n            alphas[:,s] = alpha  \n        \n        \n        return preds, alphas\n    \n    def generate_caption(self,features,max_len=200,itos=None,stoi=None):\n        # Inference part\n        # Given the image features generate the captions\n        \n        batch_size = features.size(0)\n        h, c = self.init_hidden_state(features)  # (batch_size, decoder_dim)\n        \n        alphas = []\n        \n        #starting input\n        #word = torch.tensor(stoi['<sos>']).view(1,-1).to(device)\n        word=torch.full((batch_size,1),stoi['<sos>']).to(device)\n        embeds = self.embedding(word)\n\n        \n        #captions = []\n        captions=torch.zeros((batch_size,202),dtype=torch.long).to(device)\n        captions[:,0]=word.squeeze()\n        \n        for i in range(202):\n            alpha,context = self.attention(features, h)\n            \n            \n            #store the apla score\n            #alphas.append(alpha.cpu().detach().numpy())\n            #print('embeds',embeds.shape)\n            #print('embeds[:,0]',embeds[:,0].shape)\n            #print('context',context.shape)\n            lstm_input = torch.cat((embeds[:, 0], context), dim=1)\n            h, c = self.lstm_cell(lstm_input, (h, c))\n            output = self.fcn(self.drop(h))\n            #print('output',output.shape)\n            output = output.view(batch_size,-1)\n        \n            \n            #select the word with most val\n            predicted_word_idx = output.argmax(dim=1)\n            \n            #save the generated word\n            #captions.append(predicted_word_idx.item())\n            #print('predicted_word_idx',predicted_word_idx.shape)\n            captions[:,i]=predicted_word_idx\n            \n            #end if <EOS detected>\n            #if itos[predicted_word_idx.item()] == \"<eos>\":\n            #    break\n            \n            #send generated word as the next caption\n            #embeds = self.embedding(predicted_word_idx.unsqueeze(0))\n            embeds=self.embedding(predicted_word_idx).unsqueeze(1)\n        \n        #covert the vocab idx to words and return sentence\n        #return [itos[idx] for idx in captions]\n        return captions\n    \n    \n    def init_hidden_state(self, encoder_out):\n        mean_encoder_out = encoder_out.mean(dim=1)\n        h = self.init_h(mean_encoder_out)  # (batch_size, decoder_dim)\n        c = self.init_c(mean_encoder_out)\n        return h, c\n    \n\n\nclass EncoderCNNtrain18(nn.Module):\n    def __init__(self):\n        super(EncoderCNNtrain18, self).__init__()\n        resnet = torchvision.models.resnet18()\n        #for param in resnet.parameters():\n        #    param.requires_grad_(False)\n        \n        modules = list(resnet.children())[:-2]\n        self.resnet = nn.Sequential(*modules)\n        \n\n    def forward(self, images):\n        features = self.resnet(images)                                    #(batch_size,512,8,8)\n        features = features.permute(0, 2, 3, 1)                           #(batch_size,8,8,512)\n        features = features.view(features.size(0), -1, features.size(-1)) #(batch_size,64,512)\n        #print(features.shape)\n        return features\n    \nclass EncoderDecodertrain18(nn.Module):\n    def __init__(self,embed_size, vocab_size, attention_dim,encoder_dim,decoder_dim,drop_prob=0.3):\n        super().__init__()\n        self.encoder = EncoderCNNtrain18()\n        self.decoder = DecoderRNN(\n            embed_size=embed_size,\n            vocab_size = vocab_size,\n            attention_dim=attention_dim,\n            encoder_dim=encoder_dim,\n            decoder_dim=decoder_dim\n        )\n        \n    def forward(self, images, captions):\n        features = self.encoder(images)\n        outputs = self.decoder(features, captions)\n        return outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"embed_size=200\nvocab_size = len(vocab)\nattention_dim=300\nencoder_dim=512\ndecoder_dim=300\n\nmodel = EncoderDecodertrain18(\n    embed_size=embed_size,\n    vocab_size = vocab_size,\n    attention_dim=attention_dim,\n    encoder_dim=encoder_dim,\n    decoder_dim=decoder_dim\n)\n\nMODEL_PATH='../input/model18train/modeltrain18_2'\nmodel.load_state_dict(torch.load(MODEL_PATH))\nmodel=model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = Compose([\n    #RandomHorizontalFlip(),\n    Resize((256,256), PIL.Image.BICUBIC),\n    ToTensor(),\n    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\ndataset_test=InputDatasetTest(test['file_path'].to_numpy(),transform)\ndataloader_test=DataLoader(\n    dataset=dataset_test,\n    batch_size=300,\n    shuffle=False,\n    num_workers=6)\n\ndef tensor_to_captions(ten):\n    l=ten.tolist()\n    ret=[]\n    for ls in l:\n        temp=''\n        #for i in ls[1:]:\n        for i in ls:\n            if i==stoi['<eos>'] or i==stoi['<pad>']:\n                break\n            temp=temp+itos[i]\n        ret.append(temp)\n    return ret\n\n#print out a caption to make sure model working correctly\nmodel.eval()\nitr=iter(dataloader_test)\n#print(next(itr))\nimg,idx=next(itr)\nprint(img.shape)\nprint(img[0:5].shape)\nfeatures=model.encoder(img[0:5].to(device))\ncaps = model.decoder.generate_caption(features,stoi=stoi,itos=itos)\n#caption = ''.join(caps)[:-1]\ncaptions=tensor_to_captions(caps)\nplt.imshow(img[0].numpy().transpose((1,2,0)))\nprint(captions)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\nwith torch.no_grad():\n    for i,batch in enumerate(dataloader_test):\n        img,idx=batch[0].to(device),batch[1]\n        features=model.encoder(img)\n        caps=model.decoder.generate_caption(features,stoi=stoi,itos=itos)\n        captions=tensor_to_captions(caps)\n        test['InChI'].loc[idx]=captions\n        if i%1000==0: print(i)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"output_cols = ['image_id', 'InChI']\ntest[output_cols].to_csv('submission.csv',index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test[output_cols].head()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}