{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"}],"dockerImageVersionId":30588,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# SenNet + HOA - Hacking the Human Vasculature in 3D","metadata":{}},{"cell_type":"markdown","source":"## Importing Libraries","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset\nfrom torchvision import transforms\nimport tifffile as tiff\nimport cv2\nimport torch.nn as nn\nimport albumentations as A\nimport numpy as np\nimport os\nimport time\nimport torch.nn.functional as F\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import DataLoader","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:00.111918Z","iopub.execute_input":"2023-11-27T15:56:00.112858Z","iopub.status.idle":"2023-11-27T15:56:02.637616Z","shell.execute_reply.started":"2023-11-27T15:56:00.112818Z","shell.execute_reply":"2023-11-27T15:56:02.636508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:02.639852Z","iopub.execute_input":"2023-11-27T15:56:02.640819Z","iopub.status.idle":"2023-11-27T15:56:03.658492Z","shell.execute_reply.started":"2023-11-27T15:56:02.640775Z","shell.execute_reply":"2023-11-27T15:56:03.657254Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sample Original Image & Label","metadata":{}},{"cell_type":"code","source":"base_path = '/kaggle/input/blood-vessel-segmentation/train'  \n\ndataset = 'kidney_1_dense'\n\nimages_path = os.path.join(base_path, dataset, 'images')\nlabels_path = os.path.join(base_path, dataset, 'labels')\n\nimage_files = sorted([os.path.join(images_path, f) for f in os.listdir(images_path) if f.endswith('.tif')])\nlabel_files = sorted([os.path.join(labels_path, f) for f in os.listdir(labels_path) if f.endswith('.tif')])\n\ndef show_images(images,titles= None, cmap='gray'):\n    n = len(images)\n    fig, axes = plt.subplots(1, n, figsize=(20, 10))\n    if not isinstance(axes, np.ndarray):\n        axes = [axes]\n    for idx, ax in enumerate(axes):\n        ax.imshow(images[idx], cmap=cmap)\n        if titles:\n            ax.set_title(titles[idx])\n        ax.axis('off')\n    plt.tight_layout()\n    plt.show()\n\nfirst_image = tiff.imread(image_files[981])\nfirst_label = tiff.imread(label_files[981])\n\nshow_images([first_image, first_label], titles=['First Image', 'First Label'])","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:03.660480Z","iopub.execute_input":"2023-11-27T15:56:03.660811Z","iopub.status.idle":"2023-11-27T15:56:04.604818Z","shell.execute_reply.started":"2023-11-27T15:56:03.660782Z","shell.execute_reply":"2023-11-27T15:56:04.603848Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataloader","metadata":{}},{"cell_type":"code","source":"\nclass CustomDataset(Dataset):\n    def __init__(self, image_files, mask_files, input_size=(256, 256), augmentation_transforms=None):\n        self.image_files = image_files\n        self.mask_files = mask_files\n        self.input_size = input_size\n        self.augmentation_transforms = augmentation_transforms\n\n    def __len__(self):\n        return len(self.image_files)\n\n    def __getitem__(self, idx):\n       \n        image_path = self.image_files[idx]\n        mask_path = self.mask_files[idx]\n\n        image = preprocess_image(image_path)\n        mask = preprocess_mask(mask_path)\n\n        if self.augmentation_transforms:\n            image, mask = self.augmentation_transforms(image, mask)\n\n        return image, mask","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.607423Z","iopub.execute_input":"2023-11-27T15:56:04.608250Z","iopub.status.idle":"2023-11-27T15:56:04.617387Z","shell.execute_reply.started":"2023-11-27T15:56:04.608184Z","shell.execute_reply":"2023-11-27T15:56:04.616462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preprocessing of Images","metadata":{}},{"cell_type":"code","source":"def preprocess_image(path):\n    \n    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n    img = np.tile(img[...,None],[1, 1, 3]) \n    img = img.astype('float32') \n    mx = np.max(img)\n    if mx:\n        img/=mx \n        \n    img = np.transpose(img, (2, 0, 1))\n    img_ten = torch.tensor(img)\n    return img_ten","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.619016Z","iopub.execute_input":"2023-11-27T15:56:04.619409Z","iopub.status.idle":"2023-11-27T15:56:04.637897Z","shell.execute_reply.started":"2023-11-27T15:56:04.619375Z","shell.execute_reply":"2023-11-27T15:56:04.636944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preprocessing of Masks","metadata":{}},{"cell_type":"code","source":"def preprocess_mask(path):\n    \n    msk = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n    msk = msk.astype('float32')\n    msk/=255.0\n    msk_ten = torch.tensor(msk)\n    \n    return msk_ten","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.639080Z","iopub.execute_input":"2023-11-27T15:56:04.639457Z","iopub.status.idle":"2023-11-27T15:56:04.646974Z","shell.execute_reply.started":"2023-11-27T15:56:04.639432Z","shell.execute_reply":"2023-11-27T15:56:04.646248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmentation","metadata":{}},{"cell_type":"code","source":"def augment_image(image, mask):\n    \n    image_np = image.permute(1, 2, 0).numpy()\n    mask_np = mask.numpy()\n\n    transform = A.Compose([\n        A.Resize(256,256, interpolation=cv2.INTER_NEAREST),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0),\n        A.RandomCrop(height=256, width=256, always_apply=True),\n        A.RandomBrightness(p=1),\n        A.OneOf(\n            [\n                A.Blur(blur_limit=3, p=1),\n                A.MotionBlur(blur_limit=3, p=1),\n            ],\n            p=0.9,\n        ),\n    \n    ])\n\n    augmented = transform(image=image_np, mask=mask_np)\n    augmented_image, augmented_mask = augmented['image'], augmented['mask']\n\n    augmented_image = torch.tensor(augmented_image, dtype=torch.float32).permute(2, 0, 1)\n    augmented_mask = torch.tensor(augmented_mask, dtype=torch.float32)\n\n    return augmented_image, augmented_mask","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.648038Z","iopub.execute_input":"2023-11-27T15:56:04.648316Z","iopub.status.idle":"2023-11-27T15:56:04.659406Z","shell.execute_reply.started":"2023-11-27T15:56:04.648286Z","shell.execute_reply":"2023-11-27T15:56:04.658532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_path = '/kaggle/input/blood-vessel-segmentation/train'  \n\ndataset = 'kidney_1_dense'\n\nimages_path = os.path.join(base_path, dataset, 'images')\nlabels_path = os.path.join(base_path, dataset, 'labels')\n\nimage_files = sorted([os.path.join(images_path, f) for f in os.listdir(images_path) if f.endswith('.tif')])\nlabel_files = sorted([os.path.join(labels_path, f) for f in os.listdir(labels_path) if f.endswith('.tif')])","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.660549Z","iopub.execute_input":"2023-11-27T15:56:04.660837Z","iopub.status.idle":"2023-11-27T15:56:04.686624Z","shell.execute_reply.started":"2023-11-27T15:56:04.660802Z","shell.execute_reply":"2023-11-27T15:56:04.685747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Splitting the Dataset","metadata":{}},{"cell_type":"code","source":"train_image_files, val_image_files, train_mask_files, val_mask_files = train_test_split(\n    image_files, label_files, test_size=0.2, random_state=42)","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.687703Z","iopub.execute_input":"2023-11-27T15:56:04.687967Z","iopub.status.idle":"2023-11-27T15:56:04.694380Z","shell.execute_reply.started":"2023-11-27T15:56:04.687943Z","shell.execute_reply":"2023-11-27T15:56:04.693403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = CustomDataset(train_image_files, train_mask_files, augmentation_transforms=augment_image)\nval_dataset = CustomDataset(val_image_files, val_mask_files, augmentation_transforms=augment_image)","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:04.696549Z","iopub.execute_input":"2023-11-27T15:56:04.696841Z","iopub.status.idle":"2023-11-27T15:56:04.702541Z","shell.execute_reply.started":"2023-11-27T15:56:04.696806Z","shell.execute_reply":"2023-11-27T15:56:04.701681Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)\nval_dataloader = DataLoader(val_dataset, batch_size=8, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:05.158160Z","iopub.execute_input":"2023-11-27T15:56:05.158559Z","iopub.status.idle":"2023-11-27T15:56:05.163917Z","shell.execute_reply.started":"2023-11-27T15:56:05.158530Z","shell.execute_reply":"2023-11-27T15:56:05.162988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmented Batch_1 Images & Labels Viz.","metadata":{}},{"cell_type":"code","source":"for batch_idx, (batch_images, batch_masks) in enumerate(train_dataloader):\n    print(\"Batch\", batch_idx + 1)\n    print(\"Image batch shape:\", batch_images.shape)\n    print(\"Mask batch shape:\", batch_masks.shape)\n    \n    for image, mask, image_path, mask_path in zip(batch_images, batch_masks, train_image_files, train_mask_files):\n       \n        image = image.permute((1, 2, 0)).numpy()*255.0\n        image = image.astype('uint8')\n        mask = (mask*255).numpy().astype('uint8')\n        \n        image_filename = os.path.basename(image_path)\n        mask_filename = os.path.basename(mask_path)\n        \n        plt.figure(figsize=(15, 10))\n        \n        plt.subplot(2, 4, 1)\n        plt.imshow(image, cmap='gray')\n        plt.title(f\"Original Image - {image_filename}\")\n        \n        plt.subplot(2, 4, 2)\n        plt.imshow(mask, cmap='gray')\n        plt.title(f\"Mask Image - {mask_filename}\")\n        \n        plt.tight_layout()\n        plt.show()\n    break","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:06.385486Z","iopub.execute_input":"2023-11-27T15:56:06.385867Z","iopub.status.idle":"2023-11-27T15:56:11.415488Z","shell.execute_reply.started":"2023-11-27T15:56:06.385835Z","shell.execute_reply":"2023-11-27T15:56:11.414503Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for batch_idx, (batch_images, batch_masks) in enumerate(train_dataloader):\n    print(\"Batch\", batch_idx + 1)\n    print(\"Image batch shape:\", batch_images.shape)\n    print(\"Mask batch shape:\", batch_masks.shape)","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:56:11.417521Z","iopub.execute_input":"2023-11-27T15:56:11.418367Z","iopub.status.idle":"2023-11-27T15:57:23.742521Z","shell.execute_reply.started":"2023-11-27T15:56:11.418317Z","shell.execute_reply":"2023-11-27T15:57:23.741509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for batch_idx, (batch_images, batch_masks) in enumerate(val_dataloader):\n    print(\"Batch\", batch_idx + 1)\n    print(\"Image batch shape:\", batch_images.shape)\n    print(\"Mask batch shape:\", batch_masks.shape)","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:23.744165Z","iopub.execute_input":"2023-11-27T15:57:23.744596Z","iopub.status.idle":"2023-11-27T15:57:42.278679Z","shell.execute_reply.started":"2023-11-27T15:57:23.744560Z","shell.execute_reply":"2023-11-27T15:57:42.277590Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_default_device():\n    \"\"\"Pick GPU if available, else CPU\"\"\"\n    if torch.cuda.is_available():\n        return torch.device('cuda')\n    else:\n        return torch.device('cpu')\ndevice = get_default_device()\ndevice","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.281268Z","iopub.execute_input":"2023-11-27T15:57:42.281707Z","iopub.status.idle":"2023-11-27T15:57:42.321282Z","shell.execute_reply.started":"2023-11-27T15:57:42.281676Z","shell.execute_reply":"2023-11-27T15:57:42.320249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Attention U-Net architecture","metadata":{}},{"cell_type":"markdown","source":"One thing that is worth noticing, is that I had to add some black padding around images. It is because, I am using U-Net - like architecture with learnable visual attention blocks, and spatial dimensions of signal coming from skip connections and their corresponging feature maps must fit together. That is why the pictures after padding have size (592, 576) - they can be downsampled using max pooling 4 times (the dimensions are possible to divide by 4).\n\nThe picture that helps to understand the U-Net architecture is given below. It is a encoder-decoder architecture, that means that input image is first downsampled using max pooling. This process produces feature maps, that is semantic features coming from imput image. Then, the feature maps are upsampled again, to classify each pixel in the image. To make this process more accurate, authors of the architecture decided to use skip connections: signal from feature maps is combined with input image. It helps to decide precisely, which picture should be considered positive class (in this case: blood vessel) and which one should be marked as a background.\n\nThe picture's source is the U-Net paper.\n\n![alt text](https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/u-net-architecture.png)","metadata":{}},{"cell_type":"markdown","source":"Recently, the field of Deep Learning is evolving extremely fast. New ideas and architectures are proposed frequently. One of them, that helped to improve NLP is attention. It is an idea that we can try to make our model focus on certain parts of the incoming signal (in case of text: sequences of words, in case of images: matrices with their digital representation). Computer Vision community, inspired by the work of their colleagues, have developed an idea of visual attention. For example the paper \"Learn to pay attention\" by Jetley et al describes the usage of learnable attention gate, that is additional blocks that gather signals from different layers of the network and learn where useful information from the image is. The prediction is made by concatenating the output from the attention blocks, and not by making use of the last layer of the network directly. The image from the article that describes the idea is given below: ![alt text](https://miro.medium.com/max/616/0*8r63L3yR66SVBgJR)","metadata":{}},{"cell_type":"markdown","source":"This idea was futher developed for semantic segmentation tasks by Oktay et al in the article \"Attention U-Net: Learning Where to Look for the Pancreas\". The difference is that attention blocks are now learning to filter signal from skip connections by estimating the information coming from them by comparing it with features from feature maps.\n![alt text](https://miro.medium.com/max/780/1*PdYEf-OuUWkRsm2Lfrmy6A.png)\n","metadata":{}},{"cell_type":"markdown","source":"I decided to use this idea for my dataset, because it might work quite well, and also I would learn a lot.","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super(ConvBlock, self).__init__()\n\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        x = self.conv(x)\n        return x\n\n\nclass UpConv(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super(UpConv, self).__init__()\n\n        self.up = nn.Sequential(\n            nn.Upsample(scale_factor=2),\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        x = self.up(x)\n        return x\n\n\nclass AttentionBlock(nn.Module):\n    \"\"\"Attention block with learnable parameters\"\"\"\n\n    def __init__(self, F_g, F_l, n_coefficients):\n        \"\"\"\n        :param F_g: number of feature maps (channels) in previous layer\n        :param F_l: number of feature maps in corresponding encoder layer, transferred via skip connection\n        :param n_coefficients: number of learnable multi-dimensional attention coefficients\n        \"\"\"\n        super(AttentionBlock, self).__init__()\n\n        self.W_gate = nn.Sequential(\n            nn.Conv2d(F_g, n_coefficients, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(n_coefficients)\n        )\n\n        self.W_x = nn.Sequential(\n            nn.Conv2d(F_l, n_coefficients, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(n_coefficients)\n        )\n\n        self.psi = nn.Sequential(\n            nn.Conv2d(n_coefficients, 1, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(1),\n            nn.Sigmoid()\n        )\n\n        self.relu = nn.ReLU(inplace=True)\n\n    def forward(self, gate, skip_connection):\n        \"\"\"\n        :param gate: gating signal from previous layer\n        :param skip_connection: activation from corresponding encoder layer\n        :return: output activations\n        \"\"\"\n        g1 = self.W_gate(gate)\n        x1 = self.W_x(skip_connection)\n        psi = self.relu(g1 + x1)\n        psi = self.psi(psi)\n        out = skip_connection * psi\n        return out\n\n\nclass AttentionUNet(nn.Module):\n\n    def __init__(self, img_ch=3, output_ch=1):\n        super(AttentionUNet, self).__init__()\n\n        self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        self.Conv1 = ConvBlock(img_ch, 64)\n        self.Conv2 = ConvBlock(64, 128)\n        self.Conv3 = ConvBlock(128, 256)\n        self.Conv4 = ConvBlock(256, 512)\n        self.Conv5 = ConvBlock(512, 1024)\n\n        self.Up5 = UpConv(1024, 512)\n        self.Att5 = AttentionBlock(F_g=512, F_l=512, n_coefficients=256)\n        self.UpConv5 = ConvBlock(1024, 512)\n\n        self.Up4 = UpConv(512, 256)\n        self.Att4 = AttentionBlock(F_g=256, F_l=256, n_coefficients=128)\n        self.UpConv4 = ConvBlock(512, 256)\n\n        self.Up3 = UpConv(256, 128)\n        self.Att3 = AttentionBlock(F_g=128, F_l=128, n_coefficients=64)\n        self.UpConv3 = ConvBlock(256, 128)\n\n        self.Up2 = UpConv(128, 64)\n        self.Att2 = AttentionBlock(F_g=64, F_l=64, n_coefficients=32)\n        self.UpConv2 = ConvBlock(128, 64)\n\n        self.Conv = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x):\n        \"\"\"\n        e : encoder layers\n        d : decoder layers\n        s : skip-connections from encoder layers to decoder layers\n        \"\"\"\n        e1 = self.Conv1(x)\n\n        e2 = self.MaxPool(e1)\n        e2 = self.Conv2(e2)\n\n        e3 = self.MaxPool(e2)\n        e3 = self.Conv3(e3)\n\n        e4 = self.MaxPool(e3)\n        e4 = self.Conv4(e4)\n\n        e5 = self.MaxPool(e4)\n        e5 = self.Conv5(e5)\n\n        d5 = self.Up5(e5)\n\n        s4 = self.Att5(gate=d5, skip_connection=e4)\n        d5 = torch.cat((s4, d5), dim=1) \n        d5 = self.UpConv5(d5)\n\n        d4 = self.Up4(d5)\n        s3 = self.Att4(gate=d4, skip_connection=e3)\n        d4 = torch.cat((s3, d4), dim=1)\n        d4 = self.UpConv4(d4)\n\n        d3 = self.Up3(d4)\n        s2 = self.Att3(gate=d3, skip_connection=e2)\n        d3 = torch.cat((s2, d3), dim=1)\n        d3 = self.UpConv3(d3)\n\n        d2 = self.Up2(d3)\n        s1 = self.Att2(gate=d2, skip_connection=e1)\n        d2 = torch.cat((s1, d2), dim=1)\n        d2 = self.UpConv2(d2)\n\n        out = self.Conv(d2)\n\n        return out","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.323183Z","iopub.execute_input":"2023-11-27T15:57:42.323601Z","iopub.status.idle":"2023-11-27T15:57:42.352098Z","shell.execute_reply.started":"2023-11-27T15:57:42.323565Z","shell.execute_reply":"2023-11-27T15:57:42.351235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dice Coefficient","metadata":{}},{"cell_type":"markdown","source":"For segmenation task, Dice Coefficient is often used to measure performance of the models. I havn't used it as a loss function directly, because I read that training the model with such loss function is not always stable. I decided to use Focal loss, because it has this nice property that it focuses on imbalanced positive class. It was developed for object detection, but it can be used for segmenation too. In the dataset that I am using there is quite a big class imbalance - most of the pixels don't represent blood vessels.\n\nThe picture below help us understand Dice Coefficient intuitevely. It is a 2 * |AnB| / (|A| + |B|)\n\n![alt text](https://upload.wikimedia.org/wikipedia/commons/thumb/1/1f/Intersection_of_sets_A_and_B.svg/1200px-Intersection_of_sets_A_and_B.svg.png)\n","metadata":{}},{"cell_type":"code","source":"def dice_coeff(prediction, target):\n\n    mask = np.zeros_like(prediction)\n    mask[prediction >= 0.5] = 1\n\n    inter = np.sum(mask * target)\n    union = np.sum(mask) + np.sum(target)\n    epsilon = 1e-6\n    result = np.mean(2 * inter / (union + epsilon))\n    return result","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.353159Z","iopub.execute_input":"2023-11-27T15:57:42.353480Z","iopub.status.idle":"2023-11-27T15:57:42.365860Z","shell.execute_reply.started":"2023-11-27T15:57:42.353455Z","shell.execute_reply":"2023-11-27T15:57:42.365014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Focal Loss","metadata":{}},{"cell_type":"markdown","source":"\nThe plot helped me understand the Focal Loss. It comes from the Focal Loss paper.\n![alt text](https://miro.medium.com/max/1032/1*wLf0KMIzBsXPcxsN7yBrIA.png)","metadata":{}},{"cell_type":"code","source":"class FocalLoss(nn.modules.loss._WeightedLoss):\n\n    def __init__(self, gamma=0, size_average=None, ignore_index=-100,\n                 reduce=None, balance_param=1.0):\n        super(FocalLoss, self).__init__(size_average)\n        self.gamma = gamma\n        self.size_average = size_average\n        self.ignore_index = ignore_index\n        self.balance_param = balance_param\n\n    def forward(self, input, target):\n        \n        assert len(input.shape) == len(target.shape)\n        assert input.size(0) == target.size(0)\n        assert input.size(1) == target.size(1)\n\n        logpt = - F.binary_cross_entropy_with_logits(input, target)\n        pt = torch.exp(logpt)\n\n        focal_loss = -((1 - pt) ** self.gamma) * logpt\n        balanced_focal_loss = self.balance_param * focal_loss\n        return balanced_focal_loss","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.366938Z","iopub.execute_input":"2023-11-27T15:57:42.367261Z","iopub.status.idle":"2023-11-27T15:57:42.377342Z","shell.execute_reply.started":"2023-11-27T15:57:42.367231Z","shell.execute_reply":"2023-11-27T15:57:42.376512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataloaders = {\n    'training': train_dataloader,\n    'test': val_dataloader\n}","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.378432Z","iopub.execute_input":"2023-11-27T15:57:42.378695Z","iopub.status.idle":"2023-11-27T15:57:42.386587Z","shell.execute_reply.started":"2023-11-27T15:57:42.378672Z","shell.execute_reply":"2023-11-27T15:57:42.385799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training and Testing Loop","metadata":{}},{"cell_type":"code","source":"def train_and_test(model, dataloaders, optimizer, criterion, num_epochs=100, show_images=False):\n    since = time.time()\n    best_loss = 1e10\n    \n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n\n    fieldnames = ['epoch', 'training_loss', 'test_loss', 'training_dice_coeff', 'test_dice_coeff']\n    train_epoch_losses = []\n    test_epoch_losses = []\n    for epoch in range(1, num_epochs + 1):\n\n        print(f'Epoch {epoch}/{num_epochs}')\n        print('-' * 10)\n        \n        batchsummary = {a: [0] for a in fieldnames}\n        batch_train_loss = 0.0\n        batch_test_loss = 0.0\n\n        for phase in ['training', 'test']:\n            if phase == 'training':\n                model.train()  \n            else:\n                model.eval() \n\n            for sample in iter(dataloaders[phase]):\n\n                if show_images:\n                    grid_img = make_grid(sample[0])\n                    grid_img = grid_img.permute(1, 2, 0)\n                    plt.imshow(grid_img)\n                    plt.show()\n\n                inputs = sample[0].to(device)\n                masks = sample[1].to(device)\n                \n                masks = masks.unsqueeze(1)\n                \n                optimizer.zero_grad()\n\n                with torch.set_grad_enabled(phase == 'training'):\n                    outputs = model(inputs)\n\n                    loss = criterion(outputs, masks)\n\n                    y_pred = outputs.data.cpu().numpy().ravel()\n                    y_true = masks.data.cpu().numpy().ravel()\n\n                    batchsummary[f'{phase}_dice_coeff'].append(dice_coeff(y_pred, y_true))\n\n                    if phase == 'training':\n                        loss.backward()\n                        optimizer.step()\n\n                        batch_train_loss += loss.item() * sample[0].size(0)\n\n                    else:\n                        batch_test_loss += loss.item() * sample[0].size(0)\n\n            if phase == 'training':\n                epoch_train_loss = batch_train_loss / len(dataloaders['training'])\n                train_epoch_losses.append(epoch_train_loss)\n            else:\n                epoch_test_loss = batch_test_loss / len(dataloaders['test'])\n                test_epoch_losses.append(epoch_test_loss)\n\n            batchsummary['epoch'] = epoch\n            \n            print('{} Loss: {:.4f}'.format(phase, loss))\n\n        best_loss = np.max(batchsummary['test_dice_coeff'])\n        for field in fieldnames[3:]:\n            batchsummary[field] = np.mean(batchsummary[field])\n        print(\n            f'\\t\\t\\t train_dice_coeff: {batchsummary[\"training_dice_coeff\"]}, test_dice_coeff: {batchsummary[\"test_dice_coeff\"]}')\n\n    print('Best dice coefficient: {:4f}'.format(best_loss))\n\n    return model, train_epoch_losses, test_epoch_losses","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.387648Z","iopub.execute_input":"2023-11-27T15:57:42.387952Z","iopub.status.idle":"2023-11-27T15:57:42.404037Z","shell.execute_reply.started":"2023-11-27T15:57:42.387918Z","shell.execute_reply":"2023-11-27T15:57:42.403150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = 25\n\ndef train():\n    model = AttentionUNet()\n    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n    criterion = FocalLoss(gamma=2)\n\n    trained_model, train_epoch_losses, test_epoch_losses = train_and_test(model, dataloaders, optimizer, criterion, num_epochs=epochs)\n\n    return trained_model, train_epoch_losses, test_epoch_losses\n\n\ntrained_model, train_epoch_losses, test_epoch_losses = train()","metadata":{"execution":{"iopub.status.busy":"2023-11-27T15:57:42.406216Z","iopub.execute_input":"2023-11-27T15:57:42.406490Z","iopub.status.idle":"2023-11-27T17:00:44.501903Z","shell.execute_reply.started":"2023-11-27T15:57:42.406461Z","shell.execute_reply":"2023-11-27T17:00:44.500861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(trained_model.state_dict(), 'trained_model.pth')","metadata":{"execution":{"iopub.status.busy":"2023-11-27T17:00:51.642214Z","iopub.execute_input":"2023-11-27T17:00:51.642614Z","iopub.status.idle":"2023-11-27T17:00:51.861207Z","shell.execute_reply.started":"2023-11-27T17:00:51.642582Z","shell.execute_reply":"2023-11-27T17:00:51.860312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_plot, = plt.plot(range(1, len(train_epoch_losses) + 1), train_epoch_losses, label='train loss')\ntest_plot, = plt.plot(range(1, len(test_epoch_losses) + 1), test_epoch_losses, label='test loss')\nplt.legend(handles=[train_plot, test_plot])\nplt.xlabel('Epochs')\nplt.ylabel('Loss')\nplt.title('Training and Test Loss Over Epochs')\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-11-27T17:00:54.985068Z","iopub.execute_input":"2023-11-27T17:00:54.986029Z","iopub.status.idle":"2023-11-27T17:00:55.284423Z","shell.execute_reply.started":"2023-11-27T17:00:54.985992Z","shell.execute_reply":"2023-11-27T17:00:55.283386Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_plot, = plt.plot(range(len(train_epoch_losses)-15), train_epoch_losses[15:], label='train loss')\ntest_plot, = plt.plot(range(len(test_epoch_losses)-15), test_epoch_losses[15:], label='test loss')\nplt.legend(handles=[train_plot, test_plot])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-27T17:01:29.705548Z","iopub.execute_input":"2023-11-27T17:01:29.705945Z","iopub.status.idle":"2023-11-27T17:01:29.966274Z","shell.execute_reply.started":"2023-11-27T17:01:29.705914Z","shell.execute_reply":"2023-11-27T17:01:29.965154Z"},"trusted":true},"execution_count":null,"outputs":[]}]}