{"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":"markdown","source":"# Introduction\n\"Image or Object Detection is a computer technology that processes the image and detects objects in it. People often confuse Image Detection with Image Classification. Although the difference is rather clear. If you need to classify image items, you use Classification. But if you just need to locate them, for example, find out the number of objects in the picture, you should use Image Detection.\" - Renukasoni\n\nIn this notebook we will be using the Pytorch module to learn how to do an Image detection model from scratch.","metadata":{}},{"cell_type":"code","source":"# Import frequently used modules\nfrom tqdm.notebook import tqdm\nfrom PIL import Image\nimport pandas as pd\nimport numpy as np\nimport random\nimport torch\nimport time\nimport sys\nimport os\nimport gc\n\ntqdm.pandas()\ngc.enable()\nstart_time = time.time()\ntime_limit = 7 * 3600 #7 hours time limit\n\nimport matplotlib.pyplot as plt\nfrom skimage import color\n\nt_size = 512 #This is the target size of the input image and mask\n\nrandom_seed = 123\ntorch.random.manual_seed(random_seed)\nnp.random.seed(random_seed)\nrandom.seed(random_seed)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-07-15T05:20:29.744410Z","iopub.execute_input":"2022-07-15T05:20:29.744867Z","iopub.status.idle":"2022-07-15T05:20:32.097852Z","shell.execute_reply.started":"2022-07-15T05:20:29.744775Z","shell.execute_reply":"2022-07-15T05:20:32.096099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Support Functions\n\nFunction that are custom-made and are used a lot in the following cells.","metadata":{}},{"cell_type":"code","source":"def rle_to_mask(rle, width, target_size=None):\n    if target_size == None:\n        target_size = width\n\n    rle = np.array(list(map(int, rle.split())))\n    label = np.zeros((width*width))\n    \n    for start, end in zip(rle[::2], rle[1::2]):\n        label[start:start+end] = 1\n        \n    #Convert label to image\n    label = Image.fromarray(label.reshape(width, width))\n    #Resize label\n    label = label.resize((target_size, target_size))\n    label = np.array(label).astype(float)\n    #rescale label\n    label = np.round((label - label.min())/(label.max() - label.min()))\n    \n    return label.T\n\ndef mask_to_rle(mask, orig_dim=160):\n    #Rescale image to original size\n    size = int(len(mask.flatten())**.5)\n    n = Image.fromarray(mask.reshape((size, size)))\n    n = n.resize((orig_dim, orig_dim))\n    n = np.array(n).astype(np.float32)\n    #Get pixels to flatten\n    pixels = n.T.flatten()\n    #Round the pixels using the half of the range of pixel value\n    pixels = (pixels-min(pixels) > ((max(pixels)-min(pixels))/2)).astype(int)\n    \n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0]\n    runs[1::2] -= runs[::2]\n    \n    return ' '.join(str(x) for x in runs)\n\ndef show_masked_img(img, mask, title=''):\n    #rescalse mask to 0-1 range regardless of min and max value\n    mask = (mask - mask.min())/(mask.max()-mask.min())\n    mask = torch.nan_to_num(mask) #Adress NaN incase of zerodiv error\n    #Round up the mask\n    mask = torch.round(mask)\n    mask = mask.reshape(img.shape[:2])\n    \n    fig, ax = plt.subplots(1, 3, figsize=(9, 3))\n    fig.suptitle(title, fontsize=16)\n    img, mask = img.numpy(), mask.numpy()\n    \n    ax[0].imshow(mask); ax[0].set_title('Mask')\n    ax[1].imshow(img); ax[1].set_title('Image')\n    ax[2].imshow(color.label2rgb(mask, img,\n                               bg_label=0, bg_color=(1.,1.,1.), alpha=0.25))\n    ax[2].set_title('Masked Image')\n    plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-07-15T05:20:32.100516Z","iopub.execute_input":"2022-07-15T05:20:32.101238Z","iopub.status.idle":"2022-07-15T05:20:32.151931Z","shell.execute_reply.started":"2022-07-15T05:20:32.101193Z","shell.execute_reply":"2022-07-15T05:20:32.150170Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Input Data\n\nLet's read the input train and test data and see what's in them.","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv('../input/hubmap-organ-segmentation/train.csv')\ntest_df = pd.read_csv('../input/hubmap-organ-segmentation/test.csv')\ndisplay(train_df)\ndisplay(test_df)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:32.155919Z","iopub.execute_input":"2022-07-15T05:20:32.157542Z","iopub.status.idle":"2022-07-15T05:20:32.558440Z","shell.execute_reply.started":"2022-07-15T05:20:32.157499Z","shell.execute_reply":"2022-07-15T05:20:32.557259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here, we can see that only the `organ` column can provide additional info to our model. We will therefore use only the `rle` and the `organ` columns later on.","metadata":{}},{"cell_type":"markdown","source":"Let's try and see the reversibility of our `rle_to_mask` and `mask_to_rle` functions. This should print `True` if successful.","metadata":{}},{"cell_type":"code","source":"rle = train_df.rle.iloc[0]\nl = train_df.img_width.iloc[0]\n# Test converting the rle back and forth and check for loss of data\nmask = rle_to_mask(rle, l)\nnrle = mask_to_rle(mask.flatten(), l)\nprint(\"Original is same as Reconstructed:\", rle == nrle)","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:32.560885Z","iopub.execute_input":"2022-07-15T05:20:32.561216Z","iopub.status.idle":"2022-07-15T05:20:37.114666Z","shell.execute_reply.started":"2022-07-15T05:20:32.561187Z","shell.execute_reply":"2022-07-15T05:20:37.113921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We will also create an organ mapper to convert the `organ` column into numerical IDs.","metadata":{}},{"cell_type":"code","source":"organ_mapper = {o:n for n,o in enumerate(train_df.organ.unique())}\ntrain_df['organ'] = train_df.organ.apply(lambda x: organ_mapper.get(x, 5))\norgan_mapper","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:37.115628Z","iopub.execute_input":"2022-07-15T05:20:37.116328Z","iopub.status.idle":"2022-07-15T05:20:37.128686Z","shell.execute_reply.started":"2022-07-15T05:20:37.116293Z","shell.execute_reply":"2022-07-15T05:20:37.127532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataloader\n\nNext up, we define the dataloader. This will load our data for the model for both the train and test sets. For the image, we will randomly flip it either horizontally/vertically (to save on the effort for the loader). The labels (for train) are also flipped accordingly.","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nimport torch\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Current device:\", DEVICE.upper())\n\nclass ImageDataLoader(Dataset):\n    def __init__(self, df, img_size=t_size, rotate=True, train=True):\n        self.df = df\n        self.img_size = img_size\n        self.is_train = train\n        self.rotate = rotate\n        loc = \"train\" if self.is_train else \"test\"\n        \n        self.paths = df[\"id\"].apply(\n            lambda x: f\"../input/hubmap-organ-segmentation/{loc}_images/{x}.tiff\"\n        ).values.tolist()\n        \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, idx):\n        '''Fetch image at index idx. If train, fetch the labels.'''\n        path = self.paths[idx]\n        #Read Image\n        image = Image.open(str(path))\n        image = image.resize((self.img_size, self.img_size))\n        image = np.array(image).astype(float)\n        image = (image - image.min()) / (image.max() - image.min())\n        image = torch.Tensor(image)\n        \n        #Get the organ as int\n        organ = np.array([self.df.organ.iloc[idx]])\n        \n        if self.is_train:\n            label = rle_to_mask(self.df.rle[idx], self.df.img_width[idx], self.img_size)\n            label = torch.Tensor(label)\n            if self.rotate:\n                image, label = self.flip_image(image, label)\n            label = label.reshape((1, self.img_size, self.img_size))\n            return image, organ, label\n        else:\n            if self.rotate:\n                image = self.flip_image(image)\n            return image, organ\n    \n    def flip_image(self, image, label=None):\n        '''Flips the image horizontally/vertically randomly.'''\n        flip_v = int(np.random.random()>0.5)\n        flip_h = int(np.random.random()>0.5)\n        dims = []\n        if flip_h: dims.append(0)\n        if flip_v: dims.append(1)\n        #Flip the image or label on their axis\n        img = torch.flip(image, dims)\n            \n        if self.is_train:\n            label = torch.flip(label, dims)\n            return img, label\n        else:\n            return img","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:37.130126Z","iopub.execute_input":"2022-07-15T05:20:37.130896Z","iopub.status.idle":"2022-07-15T05:20:37.149560Z","shell.execute_reply.started":"2022-07-15T05:20:37.130862Z","shell.execute_reply":"2022-07-15T05:20:37.148302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Testing rotation of image loader...\")\ntest = ImageDataLoader(train_df, rotate=True)\nfor i in range(2):\n    a, b, c = test[i]\n    print(a.shape, c.shape, b.shape)\n    print(b)\n    show_masked_img(a, c, f'Sample {i}')\n\nprint(\"Testing non-rotated image loader...\")\ntest = ImageDataLoader(train_df, rotate=False)\nfor i in range(2):\n    a, b, c = test[i]\n    print(a.shape, c.shape, b.shape)\n    print(b)\n    show_masked_img(a, c, f'Sample {i}')","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-07-15T05:20:37.151032Z","iopub.execute_input":"2022-07-15T05:20:37.151367Z","iopub.status.idle":"2022-07-15T05:20:41.206453Z","shell.execute_reply.started":"2022-07-15T05:20:37.151338Z","shell.execute_reply":"2022-07-15T05:20:41.205384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"Next, we have our model. This time, we will be referencing a U-Net model for its simplicity and high applicability to our current problem.","metadata":{}},{"cell_type":"code","source":"from torch import nn\n\nclass OrganEncoder(nn.Module):\n    '''Encodes the Organ type and adds more information to the image array.'''\n    def __init__(self, n, emb_size):\n        super().__init__()\n        self.emb_size = emb_size\n        self.emb = nn.Embedding(n, emb_size, max_norm=True)\n    \n    def forward(self, image, organ):\n        batch_size = organ.shape[0]\n        o = self.emb(organ)\n        o = image * o.reshape(batch_size, 1, 1, -1)\n        o = (o - o.min()) / (o.max() - o.min())\n        return o\n        \n\nclass StartConv(nn.Module):\n    '''Convolutional Block that received the Input Image.'''\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, (3, 3), padding='same')\n        self.conv2 = nn.Conv2d(out_c, out_c, (3, 3), padding='same')\n        self.batch_norm = nn.BatchNorm2d(out_c)\n        self.act = nn.ReLU()\n        self.do = nn.Dropout(0.25)\n    \n    def forward(self, x):\n        o = self.conv1(x)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.conv2(o)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        return self.do(o)\n\nclass DownConv(nn.Module):\n    '''Convolutional Block that Downsamples the Image.'''\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, (5, 5), padding='same')\n        self.conv2 = nn.Conv2d(out_c, out_c, (3, 3), padding='same')\n        self.batch_norm = nn.BatchNorm2d(out_c)\n        self.pool = nn.MaxPool2d(2, 2)\n        self.act = nn.ReLU()\n        self.do = nn.Dropout(0.25)\n    \n    def forward(self, x):\n        o = self.pool(x)\n        o = self.conv1(o)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.conv2(o)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        return self.do(o)\n\n    \nclass UpConv(nn.Module):\n    '''Convolutional Block that Upsamples the Image.'''\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, (5, 5), padding='same')\n        self.conv2 = nn.Conv2d(out_c, out_c, (3, 3), padding='same')\n        self.batch_norm = nn.BatchNorm2d(out_c)\n        self.tran = nn.ConvTranspose2d(out_c, out_c//2, (2, 2), (2, 2))\n        self.act = nn.ReLU()\n        self.do = nn.Dropout(0.25)\n    \n    def forward(self, x):\n        o = self.conv1(x)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.conv2(o)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.do(o)\n        o = self.tran(o)\n        return o\n\nclass LastConv(nn.Module):\n    '''Convolutional Block to output the predicted mask.'''\n    def __init__(self, in_c, out_c, n_class):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, (5, 5), padding='same')\n        self.conv2 = nn.Conv2d(out_c, out_c, (3, 3), padding='same')\n        self.conv3 = nn.Conv2d(out_c, n_class, (1, 1), padding='valid')\n        self.batch_norm = nn.BatchNorm2d(out_c)\n        self.act = nn.ReLU()\n        self.out_act = nn.Sigmoid()\n    \n    def forward(self, x):\n        o = self.conv1(x)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.conv2(o)\n        o = self.batch_norm(o)\n        o = self.act(o)\n        o = self.conv3(o)\n        return self.out_act(o)\n\nclass UNetModel(nn.Module):\n    '''UNET Based Model Architecture.'''\n    def __init__(self, n_class=1):\n        super().__init__()\n        n = 32\n        self.in_conv = StartConv(3, n)\n        self.enc = OrganEncoder(len(organ_mapper)+1, t_size)\n        self.dc1 = DownConv(n, n*2)\n        self.dc2 = DownConv(n*2, n*2**2)\n        self.dc3 = DownConv(n*2**2, n*2**3)\n        self.mp = nn.MaxPool2d(2, 2)\n        self.bridge = UpConv(n*2**3, n*2**4)\n        self.uc1 = UpConv(n*2**4, n*2**3)\n        self.uc2 = UpConv(n*2**3, n*2**2)\n        self.uc3 = UpConv(n*2**2, n*2)\n        self.lc = LastConv(n*2, n, n_class)\n    \n    def forward(self, x, organ):\n        i = self.in_conv(x.permute(0, 3, 1, 2))\n        i_enc = self.enc(i, organ)\n        i_next = torch.add(i, i_enc) / 2\n        dc1 = self.dc1(i_next)\n        dc2 = self.dc2(dc1)\n        dc3 = self.dc3(dc2)\n        b = self.mp(dc3)\n        b = self.bridge(b)\n        x3 = torch.cat([dc3, b], dim=1)\n        uc1 = self.uc1(x3)\n        x2 = torch.cat([dc2, uc1], dim=1)\n        uc2 = self.uc2(x2)\n        x1 = torch.cat([dc1, uc2], dim=1)\n        uc3 = self.uc3(x1)\n        x0 = torch.cat([i_next, uc3], dim=1)\n        o = self.lc(x0)\n        \n        return o","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:41.208794Z","iopub.execute_input":"2022-07-15T05:20:41.209254Z","iopub.status.idle":"2022-07-15T05:20:41.248168Z","shell.execute_reply.started":"2022-07-15T05:20:41.209209Z","shell.execute_reply":"2022-07-15T05:20:41.246841Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We will now test our model if it functions as expected.","metadata":{}},{"cell_type":"code","source":"%%time\nsample = torch.ones((2, t_size, t_size, 3))\nsam_or = torch.LongTensor([i%5 for i in range(2)])\nunet = UNetModel(1)\nout = unet(sample, sam_or)\nsample.shape, out.shape, unet","metadata":{"_kg_hide-output":false,"_kg_hide-input":false,"execution":{"iopub.status.busy":"2022-07-15T05:20:41.249556Z","iopub.execute_input":"2022-07-15T05:20:41.249843Z","iopub.status.idle":"2022-07-15T05:20:46.413254Z","shell.execute_reply.started":"2022-07-15T05:20:41.249817Z","shell.execute_reply":"2022-07-15T05:20:46.412118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The competition uses a Dice Coefficient Scoring Method so we will also define a custom DiceCoef class for our model evaluation.","metadata":{}},{"cell_type":"code","source":"class DiceCoef(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super().__init__()\n\n    def forward(self, y_true, y_pred, smooth=1.):\n        y_true = y_true.view(-1)\n        y_pred = y_pred.view(-1)\n        \n        #Round off y_pred\n        y_pred = torch.round((y_pred - y_pred.min()) / (y_pred.max() - y_pred.min()))\n        \n        intersection = (y_true * y_pred).sum()\n        dice = (2.0*intersection + smooth)/(y_true.sum() + y_pred.sum() + smooth)\n        \n        return dice","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:46.416341Z","iopub.execute_input":"2022-07-15T05:20:46.417007Z","iopub.status.idle":"2022-07-15T05:20:46.425511Z","shell.execute_reply.started":"2022-07-15T05:20:46.416935Z","shell.execute_reply":"2022-07-15T05:20:46.424381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"After finishing all the above items, we will now go and define our train and validation datasets. We will be splitting the train_df into 9:1 ratio (90% train, 10% valid) and feed them to our custom dataloader.","metadata":{}},{"cell_type":"code","source":"#Initialize train and validation data\nfrom sklearn.model_selection import train_test_split as tts\nfrom torch.nn.utils.rnn import pad_sequence\n\nBATCH_SIZE = 4\n\ndef data_collate(batch):\n    images, organs, labels = [], [], []\n    for data in batch:\n        images.append(data[0])\n        organs.append(data[1])\n        labels.append(data[2])\n    \n    return (torch.stack(images),\n            torch.LongTensor(np.array(organs)), \n            torch.stack(labels))\n\ndef test_collate(batch):\n    images, organs = [], []\n    for data in batch:\n        images.append(data[0])\n        organs.append(data[1])\n    \n    return (torch.stack(images),\n           torch.LongTensor(np.array(organs)))\n\ndf_train, df_valid = tts(train_df, test_size = 0.1, shuffle=True, random_state=61)\ndf_train, df_valid = df_train.reset_index(drop=True), df_valid.reset_index(drop=True)\nprint(\"Train:\", df_train.shape)\nprint(\"Validation:\", df_valid.shape)\n\ntrain_data = ImageDataLoader(df_train)\nvalid_data = ImageDataLoader(df_valid, rotate=False)\n\ntrain_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True,\n                          drop_last=True, collate_fn=data_collate, num_workers=2,\n                          prefetch_factor=max(1,BATCH_SIZE//2)\n                         )\n\nvalid_loader = DataLoader(valid_data, batch_size=BATCH_SIZE, shuffle=False,\n                          drop_last=False, collate_fn=data_collate, num_workers=2,\n                          prefetch_factor=max(1,BATCH_SIZE//2)\n                         )","metadata":{"execution":{"iopub.status.busy":"2022-07-15T05:20:46.426673Z","iopub.execute_input":"2022-07-15T05:20:46.426945Z","iopub.status.idle":"2022-07-15T05:20:46.918714Z","shell.execute_reply.started":"2022-07-15T05:20:46.426919Z","shell.execute_reply":"2022-07-15T05:20:46.917696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training\n\nWe will now enter the training stage. Here, the process doesn't deviate much from the normal train loops but I added some visualization (for reference) of the validation data upon reaching an improvement of the model. We also considered stopping the training when we reach overtime (based on our time_limit var from the top) or if we reached the limit of consecutive epochs without improvement (early stopping).\n\nThe saving of our model depends on the improvement of the Dice Coefficient (since this is basically what the contest uses for LB) so we will test and record the validation score for Dice Coefficient.","metadata":{}},{"cell_type":"code","source":"#Setup the methods for the FeatureExtractorModel\nunet_model = UNetModel(1).to(DEVICE)\nunet_optim = torch.optim.Adam(unet_model.parameters(), lr=1e-4, weight_decay=1e-6, amsgrad=False)\nunet_sched = torch.optim.lr_scheduler.CosineAnnealingLR(\n    unet_optim, T_max=len(train_loader),\n    eta_min=1e-10, last_epoch=-1\n)\nunet_loss_fn = torch.nn.BCELoss()\ndice_acc_fn = DiceCoef()\n\n#Load pretrained model from previous version of this notebook\n#unet_model.load_state_dict(torch.load(\"../input/pytorch-unet-from-scratch/unet_model.pth\",\n#                                      map_location=DEVICE))\n\n\nEPOCHS = 256\nno_improvement_limit = 32\nno_impr_count = 0\nbest_acc = -np.inf\nbest_loss = np.inf\n\nlr_list = []\nfor epoch in range(EPOCHS):\n    print(f\"Epoch [{epoch+1}/{EPOCHS}]:\")\n\n    if time.time() - start_time > time_limit:\n        print(\"Time Limit Exceeded! Ending Train Loop...\")\n        break #Avoid TimeLimitExceeded Error\n    \n    #Prime the models for training\n    unet_model.train()\n    #Reset the train losses\n    unet_train_loss = 0\n    train_dice_acc = 0\n    for images, organs, labels in tqdm(train_loader, desc='training'):\n        #Convert inputs to device (cpu or cuda)\n        images = images.to(DEVICE)\n        organs = organs.to(DEVICE)\n        labels = labels.to(DEVICE)\n        #Reset optimizer\n        unet_optim.zero_grad()\n        \n        preds = unet_model(images, organs)\n        unet_loss = unet_loss_fn(preds, labels)\n        unet_train_loss += unet_loss.item()\n        \n        #update loss and perform backpropagation on fe_model\n        unet_loss.backward()\n        unet_optim.step()\n        \n        with torch.no_grad():\n            dice_acc = dice_acc_fn(labels, preds.detach())\n            train_dice_acc = dice_acc.item()\n            \n    unet_sched.step()\n    \n    unet_train_loss /= len(train_loader)\n    train_dice_acc /= len(train_loader)\n    \n    unet_valid_loss = 0\n    valid_dice_acc = 0\n    #Next, we check for the validity of the model after some epochs\n    unet_model.eval()\n    for images, organs, labels in tqdm(valid_loader, desc='validating'):\n        #Convert inputs to device (cpu or cuda)\n        images = images.to(DEVICE)\n        organs = organs.to(DEVICE)\n        labels = labels.to(DEVICE)\n        \n        with torch.no_grad():\n            preds = unet_model(images, organs)\n            unet_val = unet_loss_fn(preds, labels)\n            unet_valid_loss += unet_val.item()\n            valid_dice_acc += dice_acc_fn(labels, preds.detach()).item()\n        #break\n    unet_valid_loss /= len(valid_loader)\n    valid_dice_acc /= len(valid_loader)\n    \n    print(f\"Epoch {epoch+1}/{EPOCHS}: \" +\n          f\"BCE_train_loss: {unet_train_loss:.4f} \" +\n          f\"BCE_valid_loss: {unet_valid_loss:.4f} \" +\n          f\"train_dice_acc: {train_dice_acc:.4f} \" +\n          f\"valid_dice_acc: {valid_dice_acc:.4f} \"\n         )\n        \n    #Finally, we'll compare the loss/acc with the saved total loss/acc\n    # and save the model on improvements\n    if valid_dice_acc >= best_acc or unet_valid_loss <= best_loss:\n        #Check the mask while training if loss improves\n        test_img = images.cpu().detach()[0]\n        test_lab = labels.cpu().detach()[0]\n        test_pred = preds.cpu().detach()[0]\n        show_masked_img(test_img, test_lab, f\"Original Epoch {epoch}\")\n        show_masked_img(test_img, test_pred, f\"Predicted Epoch {epoch}\")\n        \n        no_impr_count = 0\n        best_acc = max(valid_dice_acc, best_acc)\n        best_loss = min(unet_valid_loss, best_loss)\n        #if epoch != 0: #To avoid overwriting the notebook with baseline weights' score\n        torch.save(unet_model.state_dict(), \"unet_model.pth\")\n        print(\"Saved...\")\n        \n    else:\n        no_impr_count += 1\n        print(\"Loss did not improve...\")\n        #Uncomment to stop train upon reaching no_improvement_limit\n        if no_impr_count >= no_improvement_limit:\n            print(\"Early stopping!\")\n            break","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-07-15T05:20:46.920117Z","iopub.execute_input":"2022-07-15T05:20:46.920501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The train Dice Coefficient looks way too low though (compared to the validation Dice Coefficient), and I don't think the Dropout layer is the only thing causing this, however, we will ignore this for now. With the added visualization, we can actually see the progress of our model (though limited to only 1 sample) and how it improves each step of the way. It's pretty interesting to see and try to guess how the model is deciding what to adapt and what to drop.\n\nNext, we free a bit of memory (because this thing chugs a lot of RAM) before we enter the predicion stage.","metadata":{}},{"cell_type":"code","source":"del train_loader, valid_loader\ndel images, organs, labels, preds\ndel test_img, test_lab, test_pred\ngc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prediciton","metadata":{}},{"cell_type":"markdown","source":"Finally, here, we will be predicting the mask for our test_df. The one we have access to only has 1 item so be careful about memory issues when submitting your kernel (since the hidden test data is apparently more than the size of the train data). We will be using a batch size of 1 to avoid memory related errors though this will lead to a slower prediction process.","metadata":{}},{"cell_type":"code","source":"test_df = pd.read_csv(\"../input/hubmap-organ-segmentation/test.csv\")\ntest_df['organ'] = test_df.organ.apply(lambda x: organ_mapper.get(x, 5))\ntest_data = ImageDataLoader(test_df, rotate=False, train=False)\nBATCH_SIZE = 1\ntest_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False,\n                          drop_last=False, collate_fn=test_collate, num_workers=2,\n                          prefetch_factor=max(1, BATCH_SIZE//2)\n                        )\n\nunet_model.load_state_dict(torch.load(\"unet_model.pth\", map_location=DEVICE))\nunet_model.eval()\n\npredictions = []\nfor images, organs in tqdm(test_loader, desc='Predicting'):\n    images = images.to(DEVICE)\n    organs = organs.to(DEVICE)\n    preds = unet_model(images, organs)\n    predictions.append(preds.cpu().detach())\n        \npredictions = np.concatenate(predictions, axis=0)\norig_size = test_df.img_height.values\n\nsub = {'id':[], 'rle':[]}\nfor i, (pr, sz) in zip(test_df.id, zip(predictions, orig_size)):\n    sub['id'].append(i)\n    sub['rle'].append(mask_to_rle(pr, sz))\n\nsub = pd.DataFrame(sub)\nsub.to_csv('submission.csv', index=False)\nsub","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"After we saved our prediction, we will try and see how the predicted mask looks like...","metadata":{}},{"cell_type":"code","source":"try:\n    show_masked_img(images.cpu()[0], torch.Tensor(predictions[0]), title='Prediction Sample')\nexcept:\n    pass","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# End\n\nAnd this is the end of this notebook (for now, probably adding more later). Thank you for reading!","metadata":{}}]}