{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install monai","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:25.896799Z","iopub.execute_input":"2022-01-18T18:06:25.897247Z","iopub.status.idle":"2022-01-18T18:06:36.633616Z","shell.execute_reply.started":"2022-01-18T18:06:25.897125Z","shell.execute_reply":"2022-01-18T18:06:36.632840Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Table of contents\n1. importing libraries\n2. importing data\n3. exploratory data analysis\n4. build ML pipeline\n    - define torch.Dataset & split train/val/test data\n    - define model architecture\n    - define hyperparameters, loss, and optimization method\n5. Result analysis\n6. Submission","metadata":{"execution":{"iopub.status.busy":"2022-01-18T05:04:51.091689Z","iopub.execute_input":"2022-01-18T05:04:51.092217Z","iopub.status.idle":"2022-01-18T05:04:51.101794Z","shell.execute_reply.started":"2022-01-18T05:04:51.092176Z","shell.execute_reply":"2022-01-18T05:04:51.100978Z"}}},{"cell_type":"markdown","source":"# Import libraries","metadata":{}},{"cell_type":"code","source":"import os, glob, json, cv2, tqdm, copy\nimport pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import train_test_split\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nfrom torchvision import transforms\n\nfrom monai.networks.nets import UNet\nfrom monai.losses import Dice\n\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.axes_grid1 import make_axes_locatable","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:36.636539Z","iopub.execute_input":"2022-01-18T18:06:36.636991Z","iopub.status.idle":"2022-01-18T18:06:46.232798Z","shell.execute_reply.started":"2022-01-18T18:06:36.636947Z","shell.execute_reply":"2022-01-18T18:06:46.232022Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Import data","metadata":{}},{"cell_type":"code","source":"data_dir = '../input/sartorius-cell-instance-segmentation/'\ndf_train = pd.read_csv(os.path.join(data_dir, 'train.csv'))","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:46.234145Z","iopub.execute_input":"2022-01-18T18:06:46.234424Z","iopub.status.idle":"2022-01-18T18:06:46.778725Z","shell.execute_reply.started":"2022-01-18T18:06:46.234369Z","shell.execute_reply":"2022-01-18T18:06:46.777963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Collect annotations","metadata":{}},{"cell_type":"code","source":"dict_annotation = {}\nfor sample_id, annotation in df_train[['id', 'annotation']].values:\n    dict_annotation.setdefault(sample_id, [])\n    dict_annotation[sample_id].append(annotation)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:46.781777Z","iopub.execute_input":"2022-01-18T18:06:46.781983Z","iopub.status.idle":"2022-01-18T18:06:46.903130Z","shell.execute_reply.started":"2022-01-18T18:06:46.781957Z","shell.execute_reply":"2022-01-18T18:06:46.902478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Define function for encoding/decoding annotations","metadata":{}},{"cell_type":"code","source":"def decode_annotation(annotations, shape = (520, 704)):\n    mask = np.zeros(np.prod(shape), dtype = np.uint8)\n    for annotation in annotations:\n        annotation = annotation.split()\n        list_1s = [(int(start)-1, int(start)-1 + int(length)) for start, length in zip(annotation[0:][::2], annotation[1:][::2])]\n        for start, end in list_1s:\n            mask[start:end] = 1\n    mask = mask.reshape(shape)\n    return mask\n\ndef encode_annotation(mask):\n    pixels = mask.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0]\n    runs[1::2] -= runs[::2]\n    runs[::2] = runs[::2] + 1\n    return ' '.join(str(x) for x in runs)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:46.905426Z","iopub.execute_input":"2022-01-18T18:06:46.905660Z","iopub.status.idle":"2022-01-18T18:06:46.915844Z","shell.execute_reply.started":"2022-01-18T18:06:46.905628Z","shell.execute_reply":"2022-01-18T18:06:46.915048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Exploratory Data Analysis (EDA)\n1. sample cases\n2. distributional analysis\n    - how many unique ids are there?\n    - image shape\n    - cell types\n3. cell types\n    - how do they look different?","metadata":{}},{"cell_type":"markdown","source":"## Sample cases","metadata":{}},{"cell_type":"code","source":"df_train.info()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:46.917228Z","iopub.execute_input":"2022-01-18T18:06:46.917515Z","iopub.status.idle":"2022-01-18T18:06:46.983118Z","shell.execute_reply.started":"2022-01-18T18:06:46.917454Z","shell.execute_reply":"2022-01-18T18:06:46.982389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It seems like each data sample has meta data including id, annotation, image resolution (height and width), cell_type, plate_time, sample_data, sample_id, and elapsed_timedelta.<br>\nMost helpful information seem like to be `id`, `annotation`, `image resolution` and `cell_type`.<br>\nOther than id and annotation, the column name define what it means directly.\n1. `id`: a unique identifier for each cell sample image\n2. `annotation`: has information about the pixels where the neuron cells are located (also note that each sample image can have different number of annotation - each unique id has multiple number of rows and annotations)\n\nLet's look at how annotation looks like in detail.","metadata":{}},{"cell_type":"code","source":"sample_id = df_train['id'][0]\nsample_annotation = df_train.loc[df_train['id'] == sample_id, 'annotation'][0]\nsample_annotation","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:46.984405Z","iopub.execute_input":"2022-01-18T18:06:46.984888Z","iopub.status.idle":"2022-01-18T18:06:47.005271Z","shell.execute_reply.started":"2022-01-18T18:06:46.984850Z","shell.execute_reply":"2022-01-18T18:06:47.004641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"starts = list(map(int, sample_annotation.split(' ')[0:][::2]))\nlengths = list(map(int, sample_annotation.split(' ')[1:][::2]))\nlist_pair = [(s,l) for s,l in zip(starts, lengths)]\nlist_pair[0]","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:47.006315Z","iopub.execute_input":"2022-01-18T18:06:47.006625Z","iopub.status.idle":"2022-01-18T18:06:47.014242Z","shell.execute_reply.started":"2022-01-18T18:06:47.006589Z","shell.execute_reply":"2022-01-18T18:06:47.013550Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This means that the annotation begins from location 118145 and it has length of 6. From this, we can get (start, end) pair instead of (start,length) pair.","metadata":{}},{"cell_type":"code","source":"ends = [start+length for start, length in list_pair]\nnew_list_pair = [(s,e) for s,e in zip(starts, ends)]\nnew_list_pair[0]","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:47.015753Z","iopub.execute_input":"2022-01-18T18:06:47.016558Z","iopub.status.idle":"2022-01-18T18:06:47.023652Z","shell.execute_reply.started":"2022-01-18T18:06:47.016494Z","shell.execute_reply":"2022-01-18T18:06:47.023007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"From this new list pair, we can make masks by assigning 1s to the designated pixel locations (Note that the list of index, start and end, are the index from 1d flatten annotation mask image).","metadata":{}},{"cell_type":"code","source":"sample_image = np.array(Image.open(os.path.join(data_dir, 'train', f'{sample_id}.png')))\nsample_mask = np.zeros(np.prod(sample_image.shape))\nfor start, end in new_list_pair:\n    sample_mask[start:end] = 1","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:47.024940Z","iopub.execute_input":"2022-01-18T18:06:47.025688Z","iopub.status.idle":"2022-01-18T18:06:47.073585Z","shell.execute_reply.started":"2022-01-18T18:06:47.025652Z","shell.execute_reply":"2022-01-18T18:06:47.072921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Then we need to reshape the mask into 2d.","metadata":{}},{"cell_type":"code","source":"sample_mask = sample_mask.reshape(sample_image.shape)\n\nplt.imshow(sample_image, cmap = 'gray')\nplt.imshow(sample_mask, alpha = 0.3)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:47.074832Z","iopub.execute_input":"2022-01-18T18:06:47.075075Z","iopub.status.idle":"2022-01-18T18:06:47.421716Z","shell.execute_reply.started":"2022-01-18T18:06:47.075044Z","shell.execute_reply":"2022-01-18T18:06:47.420403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Looks like the yellowish area represent the neuron cell. Furthermore, it looks like each row has 1 segmentation for 1 cell. And because this cell image has more than 1 cell, let's bring all segmentations.","metadata":{}},{"cell_type":"code","source":"image = Image.open(os.path.join(data_dir, 'train', f'{sample_id}.png'))\nimage_np = np.array(image)\n\nannotations = dict_annotation[sample_id] # pre-defined earlier in the notebook for convenience\nmask = decode_annotation(annotations) # this code is also pre-defined earlier in the notebook for conveinence\n\nfig, axes = plt.subplots(1, 2, figsize = (30,15))\naxes[0].imshow(image_np, cmap = 'gray')\naxes[1].imshow(image_np, cmap = 'gray')\naxes[1].imshow(mask, alpha = 0.1)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:47.423321Z","iopub.execute_input":"2022-01-18T18:06:47.423832Z","iopub.status.idle":"2022-01-18T18:06:48.474410Z","shell.execute_reply.started":"2022-01-18T18:06:47.423795Z","shell.execute_reply":"2022-01-18T18:06:48.473719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Distributional analysis\n1. How many unique ids are there?","metadata":{}},{"cell_type":"code","source":"df_train['id'].unique().shape","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:48.475478Z","iopub.execute_input":"2022-01-18T18:06:48.476847Z","iopub.status.idle":"2022-01-18T18:06:48.492368Z","shell.execute_reply.started":"2022-01-18T18:06:48.476807Z","shell.execute_reply":"2022-01-18T18:06:48.491715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"2. Image shape","metadata":{}},{"cell_type":"code","source":"df_train[['height', 'width']].value_counts()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:48.495510Z","iopub.execute_input":"2022-01-18T18:06:48.495944Z","iopub.status.idle":"2022-01-18T18:06:48.512974Z","shell.execute_reply.started":"2022-01-18T18:06:48.495910Z","shell.execute_reply":"2022-01-18T18:06:48.512397Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"3. Cell types","metadata":{}},{"cell_type":"code","source":"df_train['cell_type'].unique()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:48.514389Z","iopub.execute_input":"2022-01-18T18:06:48.514930Z","iopub.status.idle":"2022-01-18T18:06:48.529346Z","shell.execute_reply.started":"2022-01-18T18:06:48.514888Z","shell.execute_reply":"2022-01-18T18:06:48.528687Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train.drop_duplicates('id')['cell_type'].value_counts()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:48.530685Z","iopub.execute_input":"2022-01-18T18:06:48.531136Z","iopub.status.idle":"2022-01-18T18:06:48.550329Z","shell.execute_reply.started":"2022-01-18T18:06:48.531101Z","shell.execute_reply":"2022-01-18T18:06:48.549605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Overall, there are 606 cell images with (520x702) image resolution. Cell images can be among cell_types cort, shsy5y, and astro.","metadata":{}},{"cell_type":"markdown","source":"## Cell types\n1. how do they look different?","metadata":{}},{"cell_type":"code","source":"for cell_type in df_train['cell_type'].unique():\n    sample_ids = df_train.loc[df_train['cell_type'] == cell_type, 'id'].unique()\n    np.random.shuffle(sample_ids)\n    fig, axes = plt.subplots(3,2, figsize = (10, 10))\n    for i in range(3):\n        sample_id = sample_ids[i]\n        image = Image.open(os.path.join(data_dir, 'train', f'{sample_id}.png'))\n        image_np = np.array(image)\n        annotations_string = df_train.loc[df_train['id'] == sample_id, 'annotation'].values\n        annotation = decode_annotation(annotations_string)\n        im0 = axes[i,0].imshow(image, cmap = 'gray')\n        divider = make_axes_locatable(axes[i,0])\n        cax = divider.append_axes('right', size='5%', pad=0.05)\n        fig.colorbar(im0, cax = cax, orientation='vertical');\n        axes[i,1].imshow(image, cmap = 'gray')\n        axes[i,1].imshow(annotation, alpha = 0.2)\n        axes[i,0].set_title(f'Raw Image (id: {sample_id})')\n        axes[i,1].set_title(f'With Annotation')\n    plt.suptitle(f'{cell_type}', fontsize = 20)\n    plt.tight_layout()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:48.551936Z","iopub.execute_input":"2022-01-18T18:06:48.552199Z","iopub.status.idle":"2022-01-18T18:06:53.057215Z","shell.execute_reply.started":"2022-01-18T18:06:48.552163Z","shell.execute_reply":"2022-01-18T18:06:53.056569Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To sum up,\n1. There are 606 input samples (images) with annotations. \n2. Each sample belongs to specific cell types among cort, shsy5y, and astro. \n3. It seems like cell images have different characteristics in cell shape and pixel intensity by each cell types\n    - this fact maybe important when I do feature engineering (image pre-processing) such as normalization in the future.\n\nIn conclusion, just by looking around the data by eyes, all samples have same image resolution and they look 'okay' enough to distinguish the background and annotation by my eyes.<br>\nFor now, we can just go ahead build a model and feed the raw image and annotation itself and see how it works.","metadata":{"execution":{"iopub.status.busy":"2022-01-18T05:09:25.138561Z","iopub.execute_input":"2022-01-18T05:09:25.138792Z","iopub.status.idle":"2022-01-18T05:09:25.148804Z","shell.execute_reply.started":"2022-01-18T05:09:25.138765Z","shell.execute_reply":"2022-01-18T05:09:25.147275Z"}}},{"cell_type":"markdown","source":"## Build dictionary for cell type","metadata":{}},{"cell_type":"code","source":"dict_cell_type = {}\nfor sample_id, cell_type in df_train[['id', 'cell_type']].values:\n    dict_cell_type[sample_id] = cell_type\n    \ncell_type_mapping = {\n    'shsy5y': 1,\n    'astro': 2,\n    'cort': 3\n}\ndict_cell_type_encoded = {sample_id: cell_type_mapping[cell_type] for sample_id, cell_type in dict_cell_type.items()}","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:53.058767Z","iopub.execute_input":"2022-01-18T18:06:53.059190Z","iopub.status.idle":"2022-01-18T18:06:53.162527Z","shell.execute_reply.started":"2022-01-18T18:06:53.059154Z","shell.execute_reply":"2022-01-18T18:06:53.161920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Build ML Pipeline\n(1) define data pipeline (2) define model architecture (3) prepare for training experiment configurations (# of epoch, batch size, optimizer, loss function...)","metadata":{}},{"cell_type":"markdown","source":"## Define data pipeline","metadata":{}},{"cell_type":"code","source":"class CellDataset(torch.utils.data.Dataset):\n    def __init__(self, \n                 list_ids,\n                 data_dir = data_dir,\n                 dict_annotation = dict_annotation,\n                 dict_cell_type_encoded = dict_cell_type_encoded,\n                 target_shape = (520, 704),\n                 image_trans = transforms.ToTensor(),\n                 target_image_res = (224,224)\n                ):\n        self.list_ids = list_ids\n        self.data_dir = data_dir\n        self.dict_annotation = dict_annotation\n        self.dict_cell_type_encoded = dict_cell_type_encoded\n        self.target_shape = target_shape\n        self.image_trans = image_trans\n        self.resize_trans = transforms.Resize(target_image_res)\n        self.target_image_res = target_image_res\n    def __len__(self):\n        return len(self.list_ids)\n    def __getitem__(self, idx):\n        cell_id = self.list_ids[idx]\n        image = Image.open(os.path.join(self.data_dir, 'train', f'{cell_id}.png'))\n        image_shape = image.size\n        image = self.resize_trans(self.image_trans(image))\n        annotations = torch.Tensor(decode_annotation(dict_annotation[cell_id], image_shape[::-1])).unsqueeze(0)\n        annotations = (self.resize_trans(annotations) > 0.5).float()\n        target = torch.zeros(4,*self.target_image_res) # number of channels are 4 including background\n        target[self.dict_cell_type_encoded[cell_id]] = annotations\n        target[0] = 1 - annotations\n        return image, target","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:53.163805Z","iopub.execute_input":"2022-01-18T18:06:53.164066Z","iopub.status.idle":"2022-01-18T18:06:53.174424Z","shell.execute_reply.started":"2022-01-18T18:06:53.164032Z","shell.execute_reply":"2022-01-18T18:06:53.173687Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"random_state = 42\n\nlist_ids = df_train['id'].unique().tolist()\ntrain_ids, val_test_ids = train_test_split(list_ids, test_size = 0.2, random_state = random_state)\nval_ids, test_ids = train_test_split(val_test_ids, test_size = 0.5, random_state = random_state)\n\ndataset = {\n    'train': CellDataset(train_ids),\n    'val': CellDataset(val_ids),\n    'test': CellDataset(test_ids),\n}\n\nfor key in dataset:\n    print(f'{key}: {len(dataset[key])}')","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:53.175343Z","iopub.execute_input":"2022-01-18T18:06:53.175603Z","iopub.status.idle":"2022-01-18T18:06:53.193724Z","shell.execute_reply.started":"2022-01-18T18:06:53.175568Z","shell.execute_reply":"2022-01-18T18:06:53.192723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"idx = 100\nimage, target = dataset['train'][idx]\ncell_id = dataset['train'].list_ids[idx]\nslice_idx = dataset['train'].dict_cell_type_encoded[cell_id]\nplt.imshow(image[0], cmap = 'gray')\nplt.imshow(target[slice_idx], alpha = 0.2)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:06:53.195231Z","iopub.execute_input":"2022-01-18T18:06:53.195553Z","iopub.status.idle":"2022-01-18T18:06:53.623926Z","shell.execute_reply.started":"2022-01-18T18:06:53.195517Z","shell.execute_reply":"2022-01-18T18:06:53.623260Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define model architecture\nI used pre-defined UNet architecture from [Monai](https://monai.io/) with 4 output channels (3 different cell types and additional 1 for background). Then, I added the final output layer to make the binary output. Overall, the final outputs of the model are 2 (one for multi-class prediction and the other for binary prediction).","metadata":{}},{"cell_type":"markdown","source":"### What is MONAI?\nQuoting from the official website...\n>\"The MONAI framework is the open-source foundation being created by Project MONAI. MONAI is a freely available, community-supported, PyTorch-based framework for deep learning in healthcare imaging. It provides domain-optimized foundational capabilities for developing healthcare imaging training workflows in a native PyTorch paradigm.\"\n\nSo this MONAI has lots of useful tools mostly for medical image analysis based on PyTorch so that we don't have to manually define model architectures, loss functions, and so on. It's open source and we can also modify and debug codes ourselves.","metadata":{}},{"cell_type":"code","source":"class UNet2Step(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.unet = UNet(\n            spatial_dims = 2,\n            in_channels = 1,\n            out_channels = 4, # 3 classes and 1 background\n            channels = (16, 32, 64, 128, 256),\n            kernel_size = 3,\n            strides = (2,2,2,2),\n            num_res_units = 2,\n            act = 'PRELU',\n            norm = 'INSTANCE',\n            dropout = 0\n        )\n        self.output_conv = nn.Conv2d(4, 1, 3, 1, 1)\n    def forward(self, x):\n        x_multi = self.unet(x) # multi-class classification\n        x = self.output_conv(x_multi)\n        return x_multi, x","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:16:50.122574Z","iopub.execute_input":"2022-01-18T18:16:50.123096Z","iopub.status.idle":"2022-01-18T18:16:50.129657Z","shell.execute_reply.started":"2022-01-18T18:16:50.123054Z","shell.execute_reply":"2022-01-18T18:16:50.128985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"net = UNet2Step()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:16:50.334388Z","iopub.execute_input":"2022-01-18T18:16:50.334685Z","iopub.status.idle":"2022-01-18T18:16:50.367532Z","shell.execute_reply.started":"2022-01-18T18:16:50.334656Z","shell.execute_reply":"2022-01-18T18:16:50.366918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training configuration\nFor loss functions, I defined 2 different losses (IoU loss for multi-class prediction and IoU loss for binary prediction). To monitor the model performance, I was monitoring (1) binary IOU score and (2) binary DICE score.","metadata":{}},{"cell_type":"code","source":"EPOCH = 50\nBATCH_SIZE = 5\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:16:51.201872Z","iopub.execute_input":"2022-01-18T18:16:51.202598Z","iopub.status.idle":"2022-01-18T18:16:51.206503Z","shell.execute_reply.started":"2022-01-18T18:16:51.202556Z","shell.execute_reply":"2022-01-18T18:16:51.205613Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# define dataloader\ndataloaders = {\n    'train': torch.utils.data.DataLoader(dataset['train'], batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)\n}\nfor split in ['val', 'test']:\n    dataloaders[split] = torch.utils.data.DataLoader(dataset[split], batch_size = BATCH_SIZE, shuffle = False, num_workers = 2)\n\n# define loss function\ncriterion_softmax = Dice(softmax = True, jaccard = True) # IOU (Jaccard index) Loss for multi-class segmentation\ncriterion_binary = Dice(sigmoid = True, jaccard = True) # IOU loss for binary segmentation and monitoring metric 1\ncriterion_binary_dice = Dice(sigmoid = True) # monitoring metric 2\n# Define optimizer\nnet = net.to(DEVICE)\noptimizer = torch.optim.Adam(net.parameters())","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:16:51.373115Z","iopub.execute_input":"2022-01-18T18:16:51.373433Z","iopub.status.idle":"2022-01-18T18:16:51.388643Z","shell.execute_reply.started":"2022-01-18T18:16:51.373401Z","shell.execute_reply":"2022-01-18T18:16:51.388017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"history = {\n    split: {\n        'loss_softmax': [], 'dice': [], 'loss': [], 'iou': []\n    }\n    for split in ['train', 'val', 'test']\n}\nfor epoch in range(EPOCH):\n    # model train\n    pbar = tqdm.tqdm(total = len(dataloaders['train']), position = 0)\n    pbar.set_description(f'Epoch (train): {epoch + 1}/{EPOCH}')\n    net.train()\n    list_loss = []\n    list_loss_softmax = []\n    list_dice = []\n    list_iou = []\n    for data, target in dataloaders['train']:\n        net.zero_grad()\n        # bring data\n        data = data.to(DEVICE)\n        target = target.to(DEVICE)\n        # inference\n        pred_multi, pred = net(data)\n        # loss and weight update\n        loss_multi = criterion_softmax(pred_multi, target)\n        loss_binary = criterion_binary(pred, target[:,1:,:,:].sum(dim=1, keepdim = True))\n        loss = loss_multi + loss_binary\n        loss.backward()\n        optimizer.step()\n        # collect metrics and history\n        list_loss.append(loss.item())\n        list_loss_softmax.append(loss_multi.item())\n        iou_score = 1-loss_binary\n        dice_score = 1-criterion_binary_dice(pred, target[:,1:,:,:].sum(dim = 1, keepdim = True))\n        list_dice.append(dice_score.item())\n        list_iou.append(dice_score.item())\n        pbar.update(1)\n        pbar.set_postfix({'Loss': f'{np.mean(list_loss):.2f}', 'Loss_softmax': f\"{np.mean(list_loss_softmax):.2f}\", 'Dice': f\"{np.mean(list_dice):.2f}\", 'IOU': f\"{np.mean(list_iou):.2f}\"})\n    history['train']['loss_softmax'].append(np.mean(list_loss_softmax))\n    history['train']['loss'].append(np.mean(list_loss))\n    history['train']['dice'].append(np.mean(list_dice))\n    history['train']['iou'].append(np.mean(list_iou))\n    pbar.close()\n    # model eval and test\n    net.eval()\n    with torch.no_grad():\n        for split in ['val', 'test']:\n            pbar = tqdm.tqdm(total = len(dataloaders[split]), position = 0)\n            pbar.set_description(f'Epoch ({split}): {epoch + 1}/{EPOCH}')\n            list_loss = []\n            list_loss_softmax = []\n            list_metrics = []\n            for data, target in dataloaders[split]:\n                # bring data\n                data = data.to(DEVICE)\n                target = target.to(DEVICE)\n                # inference\n                pred_multi, pred = net(data)\n                # loss and weight update\n                loss_multi = criterion_softmax(pred_multi, target)\n                loss_binary = criterion_binary(pred, target[:,1:,:,:].sum(dim=1, keepdim = True))\n                loss = loss_multi + loss_binary\n                # collect metrics and history\n                list_loss.append(loss.item())\n                list_loss_softmax.append(loss_multi.item())\n                iou_score = 1-loss_binary\n                dice_score = 1-criterion_binary_dice(pred, target[:,1:,:,:].sum(dim = 1, keepdim = True))\n                list_dice.append(dice_score.item())\n                list_iou.append(dice_score.item())\n                pbar.update(1)\n                pbar.set_postfix({'Loss': f'{np.mean(list_loss):.2f}', 'Loss_softmax': f\"{np.mean(list_loss_softmax):.2f}\", 'Dice': f\"{np.mean(list_dice):.2f}\", 'IOU': f\"{np.mean(list_iou):.2f}\"})\n            history[split]['loss_softmax'].append(np.mean(list_loss_softmax))\n            history[split]['loss'].append(np.mean(list_loss))\n            history[split]['dice'].append(np.mean(list_dice))\n            history[split]['iou'].append(np.mean(list_iou))\n            pbar.close()\n    # save best model\n    if history['val']['dice'][-1] == max(history['val']['dice']):\n        best_state_dict = copy.deepcopy(net.state_dict())\n        print(f\"Best model appeared with DICE score of {history['val']['dice'][-1]:.3f}\")\n    model_dict = {\n        'curr_epoch': epoch,\n        'best_state_dict': best_state_dict,\n        'curr_state_dict': net.state_dict(),\n        'history': history\n    }\n    torch.save(model_dict, 'model.pt')","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:16:51.714911Z","iopub.execute_input":"2022-01-18T18:16:51.715132Z","iopub.status.idle":"2022-01-18T18:25:38.667031Z","shell.execute_reply.started":"2022-01-18T18:16:51.715105Z","shell.execute_reply":"2022-01-18T18:25:38.666208Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_state_dict = model_dict['best_state_dict']\nnet.load_state_dict(best_state_dict)","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:25:38.669447Z","iopub.execute_input":"2022-01-18T18:25:38.670247Z","iopub.status.idle":"2022-01-18T18:25:38.681829Z","shell.execute_reply.started":"2022-01-18T18:25:38.670203Z","shell.execute_reply":"2022-01-18T18:25:38.680982Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Evaluation","metadata":{}},{"cell_type":"markdown","source":"## Training progress","metadata":{}},{"cell_type":"code","source":"history = model_dict['history']\nfor metric in ['loss', 'dice']:\n    for split in history:\n        plt.plot(history[split][metric], label = f'{split}-{metric}')\n    plt.legend()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:25:38.683298Z","iopub.execute_input":"2022-01-18T18:25:38.683737Z","iopub.status.idle":"2022-01-18T18:25:39.043146Z","shell.execute_reply.started":"2022-01-18T18:25:38.683683Z","shell.execute_reply":"2022-01-18T18:25:39.042497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Case-level sample predictions","metadata":{}},{"cell_type":"code","source":"count = 10\nfor data, target in dataset['test']:\n    count -= 1\n    if count == 0:\n        break\n    data = data.unsqueeze(0).to(DEVICE)\n    target = target.unsqueeze(0).to(DEVICE)\n    pred_multi, pred = net(data)\n    pred_sig = torch.sigmoid(pred)\n    target = target[:,1:,:,:].max(dim = 1, keepdim = True).values\n    \n    fig, axes = plt.subplots(1,2,figsize = (10,5))\n    axes[0].imshow(data[0][0].cpu().detach(), cmap = 'gray')\n    axes[0].imshow(target[0][0].cpu().detach(), alpha = 0.3)\n    axes[1].imshow(data[0][0].cpu().detach(), cmap = 'gray')\n    axes[1].imshow(pred_sig[0][0].cpu().detach(), alpha = 0.3)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-01-18T18:28:25.369116Z","iopub.execute_input":"2022-01-18T18:28:25.369369Z","iopub.status.idle":"2022-01-18T18:28:28.966904Z","shell.execute_reply.started":"2022-01-18T18:28:25.369337Z","shell.execute_reply":"2022-01-18T18:28:28.966242Z"},"trusted":true},"execution_count":null,"outputs":[]}]}