{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Introduction\n\nThis notebook implements the [Vision Transformer](https://arxiv.org/abs/2010.11929) model in order to predict the classes of bounding boxes. The code for the model has been taken from [here](https://github.com/lucidrains/vit-pytorch) and customized. \n\nI cropped the train images with the ground truth bounding boxes and applied the model on 16 * 16 flattened 2D patches. In order not to run OOM, I capped the number of patches to 2048 (in this way very large bounding boxes are ignored). \n\nA challenge are the positional encodings with this setup. As bounding boxes are of very different shapes, the positions of the flattened patches does not convey much information. I experimented with 2D relative (euclidian) distance encodings and 2D sinusodal positional encodings. However, the results do not appear to be better than with simple learnable 1D positional embeddings.\n\nThe model is trained from scratch for 12 epochs (on 80% of all bounding boxes in the training data; no pretrained weights used).\n\nImages in jpg-format have been kindly provided by [Md Awsafur Rahman](https://www.kaggle.com/awsaf49) ([dataset](https://www.kaggle.com/awsaf49/vinbigdata-original-image-dataset)) (Thanks for sharing!)."},{"metadata":{},"cell_type":"markdown","source":"### We start by quickly importing some modules."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfrom glob import glob\nimport shutil, os\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import GroupKFold\nfrom tqdm.notebook import tqdm\nimport seaborn as sns\nimport matplotlib.image as mpimg\nfrom collections import defaultdict\nimport dill\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim \nimport torch.nn.functional as F\nfrom torch.nn.utils.rnn import pad_sequence\nimport torchvision\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.utils.data as utils\nfrom torchvision import transforms\n\nfrom sklearn.feature_extraction.image import extract_patches_2d\nfrom sklearn.preprocessing import OneHotEncoder\nfrom sklearn.model_selection import train_test_split, GroupKFold\nfrom PIL import Image\nimport cv2\nimport pydicom\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\n\n!pip install einops\nfrom einops import rearrange, repeat\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### GPU use is of course recommended."},{"metadata":{"trusted":true},"cell_type":"code","source":"if torch.cuda.is_available():\n    device = torch.device(\"cuda\")\n    print(\"GPU is available\")\nelse:\n    device = torch.device(\"cpu\")\n    print(\"GPU not available, CPU used\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Loading the data and deleting \"no finding\" bounding boxes"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')\n\ntrain_df = train_df[train_df['class_id']!=14].reset_index(drop=True)\n\ntrain_df[['x_min','y_min','x_max','y_max']] = train_df[['x_min','y_min','x_max','y_max']].applymap(int)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### GroupKFold train-validation split in order to ensure that bounding boxes of the same image are not in training and validation data"},{"metadata":{"trusted":true},"cell_type":"code","source":"gkf = GroupKFold(n_splits=5)\n\ntrain_df['fold'] = -1\nfor fold, (train, val) in enumerate(gkf.split(train_df[['x_min','y_min','x_max','y_max']],train_df['class_id'],groups=train_df.image_id.to_list())): \n    train_df.loc[val,'fold']=fold\n    \nfold = 0\nval_df = train_df.loc[train_df['fold']==fold]\ntrain_df = train_df.loc[train_df['fold']!=fold]\n\ntrain_df.reset_index(drop=True, inplace=True)\nval_df.reset_index(drop=True, inplace=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Residual(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n    def forward(self, x, **kwargs):\n        return self.fn(x, **kwargs) + x\n\nclass PreNorm(nn.Module):\n    def __init__(self, dim, fn):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        self.fn = fn\n    def forward(self, x, **kwargs):\n        return self.fn(self.norm(x), **kwargs)\n\nclass FeedForward(nn.Module):\n    def __init__(self, dim, hidden_dim, dropout = 0.):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, dim),\n            nn.Dropout(dropout)\n        )\n    def forward(self, x):\n        return self.net(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Attention(nn.Module):\n    def __init__(self, dim, heads, dim_head, dropout):\n        super().__init__()\n        inner_dim = dim_head * heads\n        self.heads = heads\n        self.scale = dim ** -0.5\n\n        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)\n        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim),nn.Dropout(dropout))\n\n    def forward(self, x, mask = None):\n        b, n, _, h = *x.shape, self.heads\n        qkv = self.to_qkv(x).chunk(3, dim = -1)\n        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)\n\n        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale\n        mask_value = -torch.finfo(dots.dtype).max\n        \n        if mask is not None:\n            mask = mask[:, None, :] * mask[:, :, None]\n            mask = mask.unsqueeze(1)\n            dots.masked_fill_(mask, mask_value)\n            del mask\n\n        attn = dots.softmax(dim=-1)\n\n        out = torch.einsum('bhij,bhjd->bhid', attn, v)\n        out = rearrange(out, 'b h n d -> b n (h d)')\n        out =  self.to_out(out)\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Transformer(nn.Module):\n    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):\n        super().__init__()\n        self.layers = nn.ModuleList([])\n        for _ in range(depth):\n            self.layers.append(nn.ModuleList([\n                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),\n                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))\n            ]))\n    def forward(self, x, mask = None):\n        for attn, ff in self.layers:\n            x = attn(x, mask = mask)\n            x = ff(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Vision_Transformer(nn.Module):\n    def __init__(self, patch_size, max_patches, dim, depth, heads, mlp_dim, pool = 'cls', dim_head = 64, dropout = 0., emb_dropout = 0.):\n        \n        super().__init__()\n\n        self.patch_size = patch_size\n        self.pool = pool\n        \n        self.linear_projection = nn.Linear(patch_size**2, dim)\n        self.pos_embedding = nn.Parameter(torch.randn(1, max_patches + 1, dim))\n        self.linear_projection_emb = nn.Linear(dim//2, dim)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))\n        self.dropout = nn.Dropout(emb_dropout)\n        \n        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)\n\n        self.mlp_head = nn.Sequential(nn.LayerNorm(dim),nn.Linear(dim, 14))\n        \n    def forward(self, img):\n        \n        x = self.linear_projection(img)\n\n        cls_emb = self.cls_token.repeat(x.size(0), 1, 1)   \n        x = torch.cat((cls_emb, x), dim=1)\n\n        x = x + self.pos_embedding[:,:x.size(1),:]\n\n        x = self.dropout(x)\n        \n        mask = (x[:,:,0]==0)\n        x = self.transformer(x, mask)\n        \n        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0, :]\n        \n        x = self.mlp_head(x)\n        \n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Credits to: https://nlp.seas.harvard.edu/2018/04/03/attention.html\n\nclass NoamOpt:\n    \"Optim wrapper that implements rate.\"\n    def __init__(self, model_size, factor, warmup, optimizer):\n        self.optimizer = optimizer\n        self._step = 0\n        self.warmup = warmup\n        self.factor = factor\n        self.model_size = model_size\n        self._rate = 0\n        \n    def step(self):\n        \"Update parameters and rate\"\n        self._step += 1\n        rate = self.rate()\n        for p in self.optimizer.param_groups:\n            p['lr'] = rate\n        self._rate = rate\n        self.optimizer.step()\n        \n    def rate(self, step = None):\n        \"Implement `lrate` above\"\n        if step is None:\n            step = self._step\n        return self.factor * \\\n            (self.model_size ** (-0.5) *\n            min(step ** (-0.5), step * self.warmup ** (-1.5)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"d_model = 256 #dim\nbatch_size = 16\npatch_size = 16\nmax_patches = 2048\nnum_layers = 4 #depth\nheads = 4\nfeed_forward = 512 #mlp_dim\n\ntransformer = Vision_Transformer(patch_size, max_patches, d_model, num_layers, heads, feed_forward)\ntransformer = transformer.to(device)\n\ncriterion = nn.CrossEntropyLoss()\ncriterion =criterion.to(device)\n\nwarm_up_steps = 4000\noptimizer = NoamOpt(d_model, 1, warm_up_steps, optim.Adam(transformer.parameters(), lr=0))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Preprocessing the data"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageData(Dataset):\n    def __init__(self, labels_df, data_dir):\n        super().__init__()\n        self.labels_df = labels_df\n        self.data_dir = data_dir\n\n    def __len__(self):\n        return len(self.labels_df)\n    \n    def __getitem__(self, index):       \n        img_name = self.labels_df.image_id[index]\n        label = self.labels_df.class_id[index]\n        img_path = self.data_dir+img_name+\".jpg\"\n        img = plt.imread(img_path)\n        img = img[self.labels_df.y_min.values[index]:self.labels_df.y_max.values[index]+1,self.labels_df.x_min.values[index]:self.labels_df.x_max.values[index]+1]\n        img = torch.tensor(img, dtype=torch.long)\n        img = img.unsqueeze(0).unsqueeze(1)\n        if img.size(3)%patch_size != 0:\n            if img.size(3)%2 == 1:\n                h_padding = (patch_size - (img.size(3)%patch_size)) // 2 + 1\n            else:\n                h_padding = (patch_size - (img.size(3)%patch_size)) // 2     \n        else:\n            h_padding = 0\n        if img.size(2)%patch_size != 0:\n            if img.size(2)%2 == 1:\n                w_padding = (patch_size - (img.size(2)%patch_size)) // 2 + 1 \n            else:\n                w_padding = (patch_size - (img.size(2)%patch_size)) // 2       \n        else:\n            w_padding = 0\n        img = F.pad(img, (h_padding, h_padding, w_padding, w_padding))\n        img = img.squeeze(0).squeeze(0)\n        img = img.unfold(0, patch_size, patch_size).unfold(1, patch_size, patch_size)\n        img = img.reshape(-1, patch_size, patch_size)\n        img = img.reshape(img.size(0),-1)\n        return img, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def my_collate(batch):\n    data = [item[0] for item in batch if item[0].size(0) <= max_patches]  \n    target = [item[1] for item in batch if item[0].size(0) <= max_patches]\n    target = torch.tensor(target,dtype=torch.long)\n    data = pad_sequence(data, batch_first=True, padding_value=0)\n    data = torch.tensor(data, dtype=torch.float)\n    return [data, target]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"traindata = ImageData(train_df,'../input/vinbigdata-original-image-dataset/vinbigdata/train/')\ntrainset = DataLoader(dataset=traindata,batch_size=batch_size,collate_fn=my_collate, shuffle=True, num_workers=8, pin_memory=True)\n\nvaldata = ImageData(val_df,'../input/vinbigdata-original-image-dataset/vinbigdata/train/')\nvalset = DataLoader(dataset=valdata,batch_size=batch_size,collate_fn=my_collate,shuffle=False, num_workers=8, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training starts.."},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 12\n\ntrain_loss = []\nvalid_loss = []\nvalid_acc = []\n\nfor epoch in tqdm(range(1, epochs+1)):\n        \n    train_loss_batches = []\n    \n    ############################### Training ##################################\n    transformer.train()\n    \n    for i, (image_bboxes, target) in enumerate(tqdm(trainset)):\n        target = target.to(device)\n        image_bboxes = image_bboxes.to(device)\n \n        optimizer.optimizer.zero_grad()\n        output = transformer(image_bboxes)\n        \n        output = output.squeeze(1)\n        loss = criterion(output, target)\n        loss.backward()\n        optimizer.step()\n        \n        acc = (output.argmax(dim=1) == target).float()\n        acc_batch = acc.mean()\n\n        train_loss_batches.append(loss.item())\n        \n        if i == 0:\n            acc_train = acc\n        else:\n            acc_train = torch.cat((acc_train,acc))\n    \n    acc_train_epoch = acc_train.mean()\n    loss_train = np.mean(train_loss_batches)\n    train_loss.append(loss_train)\n    \n    ############################### Validation ################################\n    \n    transformer.eval()\n    with torch.no_grad():\n        valid_loss_batches = []\n        for i, (image_bboxes, target) in enumerate(valset):      \n            target = target.to(device)\n            image_bboxes = image_bboxes.to(device)\n\n            output = transformer(image_bboxes)\n                \n            output = output.squeeze(1)\n            loss_val = criterion(output, target)\n            acc = (output.argmax(dim=1) == target).float()\n                \n            valid_loss_batches.append(loss.item())\n                \n            if i == 0:\n                acc_val = acc\n            else:\n                acc_val = torch.cat((acc_val,acc))\n        \n        acc_val_epoch = acc_val.mean()\n        loss_val = np.mean(valid_loss_batches)\n        \n        valid_loss.append(loss_val)\n        valid_acc.append(acc_val_epoch)\n        \n    print(f\"Epoch: {epoch}............. Loss: {loss_train:.4f} - ACC: {acc_train_epoch:.4f} - Loss val: {loss_val:.4f} - ACC val: {acc_val_epoch:.4f}\")","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"print(f\"The highest Val ACC of {max(valid_acc):.4f} has been reached after {valid_acc.index(max(valid_acc))+1} epochs.\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(transformer, \"/kaggle/working/transformer_model.pt\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# credits to: https://github.com/Bjarten/early-stopping-pytorch/blob/master/MNIST_Early_Stopping_example.ipynb\n\n# visualize the loss as the network trained\nfig = plt.figure(figsize=(10,8))\nplt.plot(range(1,len(train_loss)+1),train_loss, label='Training Loss')\nplt.plot(range(1,len(valid_loss)+1),valid_loss,label='Validation Loss')\n\n# find position of lowest validation loss\nminposs = valid_loss.index(min(valid_loss))+1 \nplt.axvline(minposs, linestyle='--', color='r',label='Lowest validation loss')\n\nplt.xlabel('epochs')\nplt.ylabel('loss')\nplt.ylim(0, 3) # consistent scale\nplt.xlim(1, len(train_loss)) # consistent scale\nplt.grid(True)\nplt.legend()\nplt.tight_layout()\nplt.show()\nfig.savefig('loss_plot.png', bbox_inches='tight')","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}