{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":6927,"databundleVersionId":45059,"sourceType":"competition"}],"dockerImageVersionId":30746,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nprint(os.listdir(\"../input/carvana-image-masking-challenge/\"))\n\nimport zipfile\nimport shutil\n\nDATASET_DIR = '../input/carvana-image-masking-challenge/'\nWORKING_DIR = '/kaggle/working/'\n\nif len(os.listdir(WORKING_DIR)) <= 1:\n\n    with zipfile.ZipFile(DATASET_DIR + 'train.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n\n    with zipfile.ZipFile(DATASET_DIR + 'train_masks.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n    \n    print(\n        len(os.listdir(WORKING_DIR + 'train')),\n        len(os.listdir(WORKING_DIR + 'train_masks'))\n    )\n    \n    # Move some of the images (5088-4600) to the validation directory\n    train_dir = WORKING_DIR + 'train/'\n    val_dir = WORKING_DIR + 'val/'\n    os.mkdir(val_dir)\n    for file in sorted(os.listdir(train_dir))[4600:]:\n      shutil.move(train_dir + file, val_dir)\n    \n    # move their masks as well\n    masks_dir = WORKING_DIR + 'train_masks/'\n    val_masks_dir = WORKING_DIR + 'val_masks/'\n    os.mkdir(val_masks_dir)\n    for file in sorted(os.listdir(masks_dir))[4600:]:\n      shutil.move(masks_dir + file, val_masks_dir)\n\n    os.mkdir(WORKING_DIR + 'saved_images')","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:10.292673Z","iopub.execute_input":"2024-08-09T20:16:10.293505Z","iopub.status.idle":"2024-08-09T20:16:19.115867Z","shell.execute_reply.started":"2024-08-09T20:16:10.293475Z","shell.execute_reply":"2024-08-09T20:16:19.114747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{}},{"cell_type":"code","source":"def check_accuracy_binary(loader,model,device):\n    num_correct = 0\n    num_pixels = 0\n    dice_score = 0\n    model.eval()\n\n    with torch.no_grad():\n        for x, y in loader:\n            x = x.to(device)\n            y = y.to(device).unsqueeze(1)\n            preds = torch.sigmoid(model(x))\n            preds = (preds > 0.5).float()\n            num_correct += (preds == y).sum()\n            num_pixels += torch.numel(preds)\n            dice_score += (2 * (preds * y).sum()) / ((preds + y).sum() + 1e-8)\n\n    print(\n        f'Got {num_correct}/{num_pixels} with acc {num_correct/num_pixels*100:.2f}'\n    )\n    print(f'Dice score: {dice_score/len(loader)}')\n    model.train()\n    return dice_score/len(loader)\n\nfrom torchvision.utils import save_image\n\n\ndef save_predictions_as_imgs(loader, model, device, folder=\"saved_images/\"):\n    if not os.path.exists(folder):\n        os.makedirs(folder)\n    num_examples = 0\n    \n    model.eval()\n    for idx, (x, y) in enumerate(loader):\n        x = x.to(device=device)\n        with torch.no_grad():\n            preds = torch.sigmoid(model(x))\n            preds = (preds > 0.5).float()\n        \n        for i in range(preds.size(0)):  # Iterate over each image in the batch\n            torchvision.utils.save_image(preds[i], os.path.join(folder, f\"pred_{idx}_{i}.png\"))\n            torchvision.utils.save_image(y[i].unsqueeze(0), os.path.join(folder, f\"mask_{idx}_{i}.png\"))  # unsqueeze adds a channel dimension to the tensor\n            num_examples += 1\n            if num_examples == 10:\n                model.train()\n                return","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:19.118163Z","iopub.execute_input":"2024-08-09T20:16:19.118755Z","iopub.status.idle":"2024-08-09T20:16:19.130955Z","shell.execute_reply.started":"2024-08-09T20:16:19.118715Z","shell.execute_reply":"2024-08-09T20:16:19.129983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- Tips\n    - What is dice score?\n        - Dice score is a metric to evaluate the performance of a segmentation model. It is defined as the intersection of the predicted mask and the ground truth mask divided by the average of the number of pixels in the predicted mask and the ground truth mask.\n        - formula : $Dice = \\frac{2 \\times |X \\cap Y|}{|X| + |Y|}$","metadata":{}},{"cell_type":"markdown","source":"# Semantic Segmentation with UNet\n\n- We will build a model from scratch, and set up the data loading pipeline which will contain data augmentation using `albumentations` library.\n\n- we will train the model on the [Carvana Image Masking Challenge](https://www.kaggle.com/c/carvana-image-masking-challenge) dataset.","metadata":{}},{"cell_type":"markdown","source":"## Unet Architecture\n\n- The U-Net architecture is introduced in the paper titled [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597) by Olaf Ronneberger, Philipp Fischer, and Thomas Brox.\n    - they designed this architecture for biomedical image segmentation, but it can be used for any image segmentation task.\n\n- The U-Net architecture is symmetric, and it consists of two parts:\n    - Contracting path (Encoder)\n        - The contracting path is a typical convolutional network that consists of repeated application of convolutions, downsampling, and ReLU activations.\n    - Expansive path (Decoder)\n        - The expansive path consists of upsampling, concatenation with the corresponding cropped feature map from the contracting path, followed by convolution, and ReLU activation.\n    - Skip connections\n        - The skip connections are the concatenation of feature maps from contracting path with the corresponding feature maps in the expansive path.\n\n- The U-Net architecture is shown below:\n    - ![unet](https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/u-net-architecture.png)\n\n- the pattern they followed in the paper is:\n    - 2 3x3 convolutions with ReLU activation followed by 2x2 max pooling with stride 2 for downsampling. (this is repeated 4 times)\n    - 2 3x3 convolutions with ReLU activation for upsampling. (this is repeated 4 times)\n    - skip connections are added between the corresponding feature maps in the contracting and expansive paths.\n    - There is a 1x1 convolution at the end of the network, which is used to map each 64-component feature vector to the desired number of classes (1x1 convolution preserves the area and changes the depth).\n\n- Some things we will do that are different from the original U-Net architecture:\n    - We will use `Same` padding in the convolution layers to keep the spatial dimensions the same, unlike the original architecture which uses `Valid` padding.\n        - as a result, They did cropping in the skip connections to concatenate feature maps of the same spatial level (because the dimensions did not match)\n        - but for use the dimensions will match, so we will not need cropping and we will simply add the feature maps from the contracting path to the expansive path.\n        - The carvana winners used `Same` padding in their implementation (so it doesn't seem to affect the performance much).\n    - Instead of the transposed convolution, we might use a bilinear upsampling layer\n        - This layer will upsample the input by a factor of 2. It uses bilinear interpolation to upsample the input.\n        - in Gans, in Pro Gans, they used bilinear upsampling instead of transposed convolution because it produces better results.","metadata":{}},{"cell_type":"code","source":"import torch\nimport torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:24:22.606250Z","iopub.execute_input":"2024-08-09T20:24:22.606833Z","iopub.status.idle":"2024-08-09T20:24:22.615873Z","shell.execute_reply.started":"2024-08-09T20:24:22.606791Z","shell.execute_reply":"2024-08-09T20:24:22.614849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model\n\n- we will use the class `DoubleConv` to define a block of two 3x3 convolutions with ReLU activation. because this block is used multiple times in the network.","metadata":{}},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DoubleConv, self).__init__()\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1,bias=False), # same padding, we set bias to False because we will use a batchnorm layer after this (which cancels out the bias)\n            nn.BatchNorm2d(out_channels), # batchnorm layer, not in the original U-Net\n            nn.ReLU(inplace=True), # inplace true means it will modify the input directly, without allocating any additional output. It can sometimes slightly decrease the memory usage, but may not always be a valid operation (because the original input is destroyed)\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1,bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n    def forward(self, x):\n        return self.double_conv(x)\n    \n\nclass UNet(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(UNet, self).__init__()\n        # the channels of the Unet are hardcoded, but you can change them to your needs\n        self.features = [64, 128, 256, 512] \n        self.downs = nn.ModuleList()\n        self.ups = nn.ModuleList()\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # max pooling layer\n\n        # Encoder\n        for feature in self.features:\n            self.downs.append(DoubleConv(in_channels, feature))\n            # update in_channels for the next layer\n            in_channels = feature\n            # we will not add the max pooling layer here, because we want to save the outputs before the max pooling layer (for the skip connections)\n\n        # the bottleneck layer (which is the bottom of the U)\n        # this is a single level with no symmetrical upsampling (that is why we did not include 1024 in the features list) and we did it outside the loops\n        self.bottleneck = DoubleConv(self.features[-1], self.features[-1]*2) \n\n        # Decoder\n        for feature in reversed(self.features):\n            # the shift introduced because of the bottleneck layer will cause each transpose convolution to take 2*feature as input (check the diagram)\n            self.ups.append(nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2))  # kernel size 2, stride will double the width and height\n            # the double conv will take double the number of features as input because of the concatenation of the skip connection\n            self.ups.append(DoubleConv(feature*2, feature)) \n\n        # the final layer, which is a 1x1 convolution\n        self.final_layer = nn.Conv2d(self.features[0], out_channels, kernel_size=1)\n\n    def forward(self, x):\n        skip_connections = []\n\n        # Encoder\n        for level in self.downs:\n            x = level(x)\n            skip_connections.append(x)\n            x = self.pool(x)\n\n        # Bottleneck\n        x = self.bottleneck(x)\n\n        # reverse the skip connections (because we want to concatenate the layers in the reverse order)\n        skip_connections = skip_connections[::-1]\n\n        # Decoder\n        for i in range(0, len(self.ups), 2): # iterate with a step of 2 because we have 2 layers for each level\n            x = self.ups[i](x)\n            # add the skip connection\n            skip_connection = skip_connections[i//2] # i//2 because we move here with a step of 2, but we want to move with a step of 1 in the skip connections list\n            \n            # Defensive Act: check if the dimensions of the skip connection and the x are not the same (that might result when choosing an input image which size is not divisible by 2^4), which will lead to an odd number of pixels, which max pool will floor to the nearest integer, and the transpose convolution will double the width and height to that integer, which will cause the dimensions to be different by 1 pixel\n            if x.shape != skip_connection.shape:\n                # interpolate x to the size of the skip connection\n                x = F.interpolate(x, size=skip_connection.shape[2:], mode='bilinear', align_corners=True)\n\n            # concatenate the skip connection\n            x = torch.cat((skip_connection, x), dim=1) # dim=1 because we want to concatenate along the channels\n            x = self.ups[i+1](x) # the double conv layer\n\n        # Final layer\n        x = self.final_layer(x)\n\n        return x\n    \n\n# test the model\ndummy = torch.randn((3, 3, 160, 160))                        \nmodel = UNet(in_channels=3, out_channels=1)\nmodel(dummy).shape","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:19.551858Z","iopub.execute_input":"2024-08-09T20:16:19.552234Z","iopub.status.idle":"2024-08-09T20:16:21.030750Z","shell.execute_reply.started":"2024-08-09T20:16:19.552206Z","shell.execute_reply":"2024-08-09T20:16:21.029700Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- Tips\n    - We used something called ModuleList to store the layers in the network. this is different from a regular python list because it registers the layers in the network, so they can be used in the forward pass.\n        - long story short, it is compatible with the PyTorch model functions and used when we want to store layers in a list.\n\n    - ConvTranspose2d is used for upsampling in the expansive path. it takes the following arguments:\n\n    - We had to choose an input size that is divisible by 16 (because we will divide by 2 four times in the contracting path).\n        - so 2*2*2*2 = 16\n        - the reason for that is that if it is not divisible by 16, we will have an odd number of pixels, and the max pool layer will floor it when it reduces the size by half, leading to a mismatch in the dimensions when we add the skip connections later\n        - other thing we can do is to resize the skip connections or the upsampled feature maps to match the dimensions. (which we did above)\n\n        ","metadata":{}},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"import os\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\n\n\nclass CarvanaDataset(Dataset):\n    def __init__(self, image_dir, mask_dir, transform=None):\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n        self.images = os.listdir(image_dir)\n\n    def __len__(self):\n        return len(self.images)\n    \n    def __getitem__(self, index):\n        img_path = os.path.join(self.image_dir, self.images[index])\n        mask_path = os.path.join(self.mask_dir, self.images[index].replace(\".jpg\", \"_mask.gif\"))\n        # read the image and the mask (label), and store them in the form of numpy arrays (for the albumentations library)\n        image = np.array(Image.open(img_path).convert(\"RGB\")) # we might not need to do this because the images are loaded as RGB by default\n        mask = np.array(Image.open(mask_path).convert(\"L\"), dtype=np.float32) # convert(\"L\") will convert the image to grayscale\n        # binarize the mask\n        mask[mask == 255.0] = 1.0\n\n        # apply the transformations if they exist\n        if self.transform is not None:\n            augmentations = self.transform(image=image, mask=mask)\n            image = augmentations[\"image\"]\n            mask = augmentations[\"mask\"]\n            \n        return image, mask ","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.032891Z","iopub.execute_input":"2024-08-09T20:16:21.033554Z","iopub.status.idle":"2024-08-09T20:16:21.042986Z","shell.execute_reply.started":"2024-08-09T20:16:21.033519Z","shell.execute_reply":"2024-08-09T20:16:21.042337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\ntrain_image_dir = '/kaggle/working/train'\ntrain_mask_dir = '/kaggle/working/train_masks'\nval_image_dir = '/kaggle/working/val'\nval_mask_dir = '/kaggle/working/val_masks'\n## Dataset hyperparameters\nbatch_size = 16\nimage_height = 320\nimage_width = 480\npin_memory = True\n\ntrain_transform = A.Compose(\n    [\n        A.Resize(height=image_height, width=image_width),\n        A.Rotate(limit=35, p=1.0),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.1),\n        # this will only divide by 255 (since mean = 0 and std = 1)\n        A.Normalize(\n            mean=[0.0, 0.0, 0.0],\n            std=[1.0, 1.0, 1.0],\n            max_pixel_value=255.0,\n        ),\n        ToTensorV2(),\n    ],\n)\n\n# for the validation transforms, we will only resize and normalize without any augmentations\nval_transform = A.Compose(\n    [\n        A.Resize(height=image_height, width=image_width),\n        A.Normalize(\n            mean=[0.0, 0.0, 0.0],\n            std=[1.0, 1.0, 1.0],\n            max_pixel_value=255.0,\n        ),\n        ToTensorV2(),\n    ],\n)\n\n\ntrain_dataset = CarvanaDataset(image_dir=train_image_dir,mask_dir=train_mask_dir, transform=train_transform)\nval_dataset = CarvanaDataset(image_dir=val_image_dir,mask_dir=val_mask_dir, transform=val_transform)","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.044488Z","iopub.execute_input":"2024-08-09T20:16:21.045141Z","iopub.status.idle":"2024-08-09T20:16:21.218476Z","shell.execute_reply.started":"2024-08-09T20:16:21.045109Z","shell.execute_reply":"2024-08-09T20:16:21.217209Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_dataset), len(val_dataset)","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.221263Z","iopub.execute_input":"2024-08-09T20:16:21.222093Z","iopub.status.idle":"2024-08-09T20:16:21.235543Z","shell.execute_reply.started":"2024-08-09T20:16:21.222047Z","shell.execute_reply":"2024-08-09T20:16:21.234305Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset[0][0].shape, train_dataset[0][1].shape","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.237287Z","iopub.execute_input":"2024-08-09T20:16:21.239905Z","iopub.status.idle":"2024-08-09T20:16:21.393491Z","shell.execute_reply.started":"2024-08-09T20:16:21.239864Z","shell.execute_reply":"2024-08-09T20:16:21.392070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=pin_memory)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=pin_memory)","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.395673Z","iopub.execute_input":"2024-08-09T20:16:21.397081Z","iopub.status.idle":"2024-08-09T20:16:21.406248Z","shell.execute_reply.started":"2024-08-09T20:16:21.397030Z","shell.execute_reply":"2024-08-09T20:16:21.404370Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for images, masks in train_loader:\n    print(images.shape, masks.shape)\n    break","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:21.409502Z","iopub.execute_input":"2024-08-09T20:16:21.410298Z","iopub.status.idle":"2024-08-09T20:16:22.172492Z","shell.execute_reply.started":"2024-08-09T20:16:21.410257Z","shell.execute_reply":"2024-08-09T20:16:22.171570Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initializations","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nbest_accuracy = 0\n\n## Model hyperparameters\nin_channels = 3\nout_channels = 1\nlearning_rate = 1e-4\nnum_epochs = 20\n\nmodel = UNet(in_channels=in_channels, out_channels=out_channels).to(device)\n# criterion = nn.CrossEntropyLoss() # we will use this loss function if we have multiple classes (out_channels > 1)\ncriterion = nn.BCEWithLogitsLoss() # binary cross entropy with logits loss (it expects the logits, and it will apply the sigmoid function by itself), the sigmoid is applied because we have only 1 channel, and each pixel will have a value between 0 and 1 that represents the probability of that pixel being class 1\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n\n# scalar\nscaler = torch.cuda.amp.GradScaler() # this will help us to use mixed precision training\n\nuse_scheduler = True\nif use_scheduler:\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=1e-6, last_epoch=-1)\n\n\n# if we want to load a model and continue training\nload_model = False\nif load_model:\n    checkpoint = torch.load(\"model.pth\")\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n    best_accuracy = checkpoint[\"best_accuracy\"]\n    print(\"=> Loaded model with accuracy {:.2f}\".format(best_accuracy))","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:22.174515Z","iopub.execute_input":"2024-08-09T20:16:22.175157Z","iopub.status.idle":"2024-08-09T20:16:22.679450Z","shell.execute_reply.started":"2024-08-09T20:16:22.175124Z","shell.execute_reply":"2024-08-09T20:16:22.677792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train","metadata":{}},{"cell_type":"code","source":"for epoch in range(num_epochs):\n    ## Training phase \n    model.train()\n    tk0 = tqdm(train_loader, total=len(train_loader), desc=f\"Epoch {epoch+1} Training\")\n\n    train_loss = 0\n    train_exapmles = 0\n\n    # loop on the train loader\n    for batch_idx, (images, masks) in enumerate(tk0):\n        images = images.to(device)\n        masks = masks.float().unsqueeze(1).to(device) # add a channel dimension to the mask (since it is a single channel image)\n\n        \n        # we will use Float16 training to speed up the training process\n        with torch.cuda.amp.autocast():\n            # forward pass\n            preds = model(images)\n\n            # calculate the loss\n            loss = criterion(preds, masks)\n            train_loss += loss.item()\n            train_exapmles += images.size(0)\n        \n        \n        # backpropagation\n        optimizer.zero_grad()\n        scaler.scale(loss).backward() # scale the loss to avoid underflow or overflow\n        \n        # update the weights\n        scaler.step(optimizer)\n        scaler.update()\n\n        # update the progress bar\n        tk0.set_postfix(loss=(train_loss/train_exapmles))\n\n    ## Validation phase\n    model.eval()\n    tk1 = tqdm(val_loader, total=len(val_loader), desc=f\"Epoch {epoch+1} Validation\")\n\n    val_loss = 0\n    val_examples = 0\n\n    with torch.no_grad():\n        for batch_idx, (images, masks) in enumerate(tk1):\n            images = images.to(device)\n            masks = masks.float().unsqueeze(1).to(device)\n\n            # forward pass\n            preds = model(images)\n\n            # calculate the loss\n            loss = criterion(preds, masks)\n            val_loss += loss.item()\n            val_examples += images.size(0)\n\n            # update the progress bar\n            tk1.set_postfix(loss=(val_loss/val_examples))\n\n    # save the model if the accuracy is improved\n    accuracy = check_accuracy_binary(val_loader, model, device)\n    if accuracy > best_accuracy:\n        best_accuracy = accuracy\n        checkpoint = {\n            \"state_dict\": model.state_dict(),\n            \"optimizer\": optimizer.state_dict(),\n            \"best_accuracy\": best_accuracy,\n        }\n        torch.save(checkpoint, \"model.pth\")\n\n    print(f\"Epoch {epoch+1}, train loss: {train_loss/train_exapmles}, val loss: {val_loss/val_examples}, val accuracy: {accuracy}\")\n\n    # save the predictions as images every 5 epochs\n    if epoch % 5 == 0:\n        save_predictions_as_imgs(val_loader, model, device, folder=\"saved_images/\")\n    \n    if use_scheduler:\n        scheduler.step()        ","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:16:22.727890Z","iopub.execute_input":"2024-08-09T20:16:22.728269Z","iopub.status.idle":"2024-08-09T20:23:53.974180Z","shell.execute_reply.started":"2024-08-09T20:16:22.728245Z","shell.execute_reply":"2024-08-09T20:23:53.972992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submission","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport time\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[:-1:2]\n    \n    return ' '.join(str(x) for x in runs)","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:24:38.981358Z","iopub.execute_input":"2024-08-09T20:24:38.981744Z","iopub.status.idle":"2024-08-09T20:24:39.313122Z","shell.execute_reply.started":"2024-08-09T20:24:38.981713Z","shell.execute_reply":"2024-08-09T20:24:39.312259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 'test' not in os.listdir(WORKING_DIR):\n    if os.path.isfile(WORKING_DIR + 'submission.csv'):\n        os.remove(WORKING_DIR + 'submission.csv')\n    with zipfile.ZipFile(DATASET_DIR + 'test.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\nif 'test_images' not in os.listdir(WORKING_DIR):\n    os.mkdir(WORKING_DIR + 'test_images')","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:24:40.068394Z","iopub.execute_input":"2024-08-09T20:24:40.069002Z","iopub.status.idle":"2024-08-09T20:27:11.177914Z","shell.execute_reply.started":"2024-08-09T20:24:40.068958Z","shell.execute_reply":"2024-08-09T20:27:11.176938Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TEST_DIR = WORKING_DIR + 'test'\nTHRESHOLD = 0.5\nimport torchvision.transforms.functional as TF\n\n# Dataset\nclass CarvanaTestDataset(Dataset):\n    def __init__(self, image_dir, transform=None):\n        self.image_dir = image_dir\n        self.transform = transform\n        self.images = sorted(os.listdir(image_dir))\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_name = self.images[index]\n        img_path = os.path.join(self.image_dir, self.images[index])\n        image = np.array(Image.open(img_path).convert('RGB'))\n\n        if self.transform is not None:\n            augmentations = self.transform(image=image)\n            image = augmentations['image']\n\n        return img_name, image\n\n\ntest_transform = A.Compose(\n    [\n        A.Resize(height=image_height, width=image_width),\n        A.Normalize(\n            mean=[0.0, 0.0, 0.0],\n            std=[1.0, 1.0, 1.0],\n            max_pixel_value=255.0,\n        ),\n        ToTensorV2(), \n    ]\n)\n   \ntest_set = CarvanaTestDataset(\n    image_dir=TEST_DIR,\n    transform=test_transform\n)    \n\n\ntest_loader = DataLoader(\n    test_set, batch_size=batch_size, shuffle=False\n)\n    \n# Model\ncheckpoint = torch.load(WORKING_DIR + 'model.pth')\n\nmodel = UNet(in_channels=3, out_channels=1).to(device)\nmodel.load_state_dict(checkpoint['state_dict'])\n\nmodel.eval()\n\n# Predictions\nall_predictions = []\nfor img_names, x in tqdm(test_loader):\n    x = x.to(device)\n    with torch.no_grad():\n        preds = torch.sigmoid(model(x))\n        preds = (preds > THRESHOLD).float()   \n    preds = TF.resize(\n        preds, size=(1280, 1918), interpolation=TF.InterpolationMode.NEAREST\n    )\n    \n    # Encoding\n    for idx in range(len(img_names)):\n        encoding = rle_encode(preds[idx].squeeze().cpu())\n        all_predictions.append([img_names[idx], encoding])","metadata":{"execution":{"iopub.status.busy":"2024-08-09T20:27:27.044114Z","iopub.execute_input":"2024-08-09T20:27:27.044984Z","iopub.status.idle":"2024-08-09T20:27:39.182943Z","shell.execute_reply.started":"2024-08-09T20:27:27.044949Z","shell.execute_reply":"2024-08-09T20:27:39.181820Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree(WORKING_DIR + 'test')\nsub = pd.DataFrame(all_predictions)\nsub.columns = ['img', 'rle_mask']\nsub.to_csv(os.path.join(WORKING_DIR, 'submission.csv'), index=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom IPython.display import FileLink\n\nos.chdir('/kaggle/working')\nFileLink('submission.csv')","metadata":{},"execution_count":null,"outputs":[]}]}