{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch Custom Resnet - Pneumonia Detection"},{"metadata":{"trusted":false},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport glob\nimport random\nimport matplotlib.pyplot as plt\nfrom skimage.transform import resize\nfrom tqdm.notebook import tqdm\nfrom pydicom import dcmread\nimport pydicom\nimport csv\nimport os","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"from tensorflow import keras\n\nimport torch\nfrom torch.utils import data\nimport torch.optim as optim\nfrom torchvision.models import resnet50\n\nimport torch.nn as nn\nimport torch.nn.functional as F","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Retrieve the pneumonia locations for each dicom file"},{"metadata":{"trusted":false},"cell_type":"code","source":"# empty dictionary\npneumonia_locations = {}\n# load table\nwith open(os.path.join('../input/rsna-pneumonia-detection-challenge/stage_2_train_labels.csv'), mode='r') as infile:\n    # open reader\n    reader = csv.reader(infile)\n    # skip header\n    next(reader, None)\n    # loop through rows\n    for rows in reader:\n        # retrieve information\n        filename = rows[0]\n        location = rows[1:5]\n        pneumonia = rows[5]\n        # if row contains pneumonia add label to dictionary\n        # which contains a list of pneumonia locations per filename\n        if pneumonia == '1':\n            # convert string to float to int\n            location = [int(float(i)) for i in location]\n            # save pneumonia location in dictionary\n            if filename in pneumonia_locations:\n                pneumonia_locations[filename].append(location)\n            else:\n                pneumonia_locations[filename] = [location]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Split Train + Validation dataset"},{"metadata":{"trusted":false},"cell_type":"code","source":"# load and shuffle filenames\nfolder = '../input/rsna-pneumonia-detection-challenge/stage_2_train_images'\nfilenames = os.listdir(folder)\nrandom.shuffle(filenames)\n# split into train and validation filenames\nn_valid_samples = 2560\ntrain_filenames = filenames[n_valid_samples:]\nvalid_filenames = filenames[:n_valid_samples]\nprint('n train samples', len(train_filenames))\nprint('n valid samples', len(valid_filenames))\nn_train_samples = len(filenames) - n_valid_samples","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset class for each dicom file split in batches"},{"metadata":{"trusted":false},"cell_type":"code","source":"class Dataset(data.Dataset):\n    def __init__(self, device, folder, filenames, pneumonia_locations=None, batch_size=32, image_size=320, shuffle=True, augment=False, predict=False):\n        self.device = device\n        self.folder = folder\n        self.filenames = filenames\n        self.pneumonia_locations = pneumonia_locations\n        self.batch_size = batch_size\n        self.image_size = image_size\n        self.shuffle = shuffle\n        self.augment = augment\n        self.predict = predict\n        self.on_epoch_end()s\n\n    def reshape(self, imgs, msks):\n        # reshape for pytorch\n        imgs = imgs.permute(0, 3, 1, 2)\n        msks = msks.permute(0, 3, 1, 2)\n        return imgs, msks\n\n    def __load__(self, filename):\n        # load dicom file as numpy array\n        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n        # create empty mask\n        msk = np.zeros(img.shape)\n        # get filename without extension\n        filename = filename.split('.')[0]\n        # if image contains pneumonia\n        if filename in pneumonia_locations:\n            # loop through pneumonia\n            for location in pneumonia_locations[filename]:\n                # add 1's at the location of the pneumonia\n                x, y, w, h = location\n                msk[y:y+h, x:x+w] = 1\n        # if augment then horizontal flip half the time\n        if self.augment and random.random() > 0.5:\n            img = np.fliplr(img)\n            msk = np.fliplr(msk)\n        # resize both image and mask\n        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n        msk = resize(msk, (self.image_size, self.image_size), mode='reflect') > 0.5\n        # add trailing channel dimension\n        img = np.expand_dims(img, -1)\n        msk = np.expand_dims(msk, -1)\n        return img, msk\n    \n    def __loadpredict__(self, filename):\n        # load dicom file as numpy array\n        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n        # resize image\n        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n        # add trailing channel dimension\n        img = np.expand_dims(img, -1)\n        return img\n        \n    def __getitem__(self, index):\n        # select batch\n        filenames = self.filenames[index*self.batch_size:(index+1)*self.batch_size]\n        # predict mode: return images and filenames\n        if self.predict:\n            # load files\n            imgs = [self.__loadpredict__(filename) for filename in filenames]\n            # create numpy batch\n            imgs = np.array(imgs)\n            return imgs, filenames\n        # train mode: return images and masks\n        else:\n            # load files\n            items = [self.__load__(filename) for filename in filenames]\n            # unzip images and masks\n            imgs, msks = zip(*items)\n            # create numpy batch\n            imgs = torch.tensor(imgs).to(device).type(torch.cuda.FloatTensor)\n            msks = torch.tensor(msks).to(device).type(torch.cuda.FloatTensor)\n            imgs, msks = self.reshape(imgs, msks)\n            return imgs, msks\n        \n    def on_epoch_end(self):\n        if self.shuffle:\n            random.shuffle(self.filenames)\n        \n    def __len__(self):\n        if self.predict:\n            # return everything\n            return int(np.ceil(len(self.filenames) / self.batch_size))\n        else:\n            # return full batches only\n            return int(len(self.filenames) / self.batch_size)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"BATCH_SIZE = 8\nIMAGE_SIZE = 320","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"folder = '../input/rsna-pneumonia-detection-challenge/stage_2_train_images'\ntrain_gen = Dataset(device, folder, train_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=True, predict=False)\nvalid_gen = Dataset(device, folder, valid_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=False, predict=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Pytorch Custom ResNet Model"},{"metadata":{"trusted":false},"cell_type":"code","source":"class conv_block(nn.Module):\n    \"\"\"\n    Define the [convolution - batch normalization - activation] block \n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True,\n                 bn_momentum=0.9, alpha_leaky=0.03):\n        super(conv_block, self).__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n                              stride=stride, padding=padding, bias=bias)\n        self.bn = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=bn_momentum)\n        self.activ = nn.LeakyReLU(negative_slope=alpha_leaky)\n\n    def forward(self, x):\n        return self.activ(self.bn(self.conv(x)))\n    \n\nclass conv_t_block(nn.Module):\n    \"\"\"\n    Define the [convolution_transpose - batch normalization - activation] block \n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, output_size=None, kernel_size=3, bias=True,\n                 bn_momentum=0.9, alpha_leaky=0.03):\n        super(conv_t_block, self).__init__()\n        self.conv_t = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=2, padding=1, \n                                         bias=bias)\n        self.bn = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=bn_momentum)\n        self.activ = nn.LeakyReLU(negative_slope=alpha_leaky)\n\n    def forward(self, x, output_size):\n        return self.activ(self.bn(self.conv_t(x, output_size=output_size)))   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class Net(nn.Module):\n    '''\n    Custom ResNet module\n    '''\n\n    def __init__(self):\n        super(Net, self).__init__()\n        \n        self.down_1 = nn.Sequential(conv_block(in_channels=1, out_channels=64), conv_block(in_channels=64, out_channels=64))\n        self.down_2 = nn.Sequential(conv_block(in_channels=64, out_channels=128), conv_block(in_channels=128, out_channels=128))\n        self.down_3 = nn.Sequential(conv_block(in_channels=128, out_channels=256), conv_block(in_channels=256, out_channels=256))\n        self.down_4 = nn.Sequential(conv_block(in_channels=256, out_channels=512), conv_block(in_channels=512, out_channels=512))\n        self.down_5 = nn.Sequential(conv_block(in_channels=512, out_channels=512), conv_block(in_channels=512, out_channels=512))\n\n        self.middle = nn.Sequential(conv_block(in_channels=512, out_channels=512), conv_block(in_channels=512, out_channels=512))\n        self.middle_t = conv_t_block(in_channels=512, out_channels=256)\n\n        self.up_5 = nn.Sequential(conv_block(in_channels=768, out_channels=512), conv_block(in_channels=512, out_channels=512))\n        self.up_5_t = conv_t_block(in_channels=512, out_channels=256)\n        self.up_4 = nn.Sequential(conv_block(in_channels=768, out_channels=512), conv_block(in_channels=512, out_channels=512))\n        self.up_4_t = conv_t_block(in_channels=512, out_channels=128)\n        self.up_3 = nn.Sequential(conv_block(in_channels=384, out_channels=256), conv_block(in_channels=256, out_channels=256))\n        self.up_3_t = conv_t_block(in_channels=256, out_channels=64)\n        self.up_2 = nn.Sequential(conv_block(in_channels=192, out_channels=128), conv_block(in_channels=128, out_channels=128))\n        self.up_2_t = conv_t_block(in_channels=128, out_channels=32)\n        self.up_1 = nn.Sequential(conv_block(in_channels=96, out_channels=64), conv_block(in_channels=64, out_channels=1))\n        \n    def forward(self, x):\n        down1 = self.down_1(x)\n        out = F.max_pool2d(down1, kernel_size=2, stride=2)\n\n        down2 = self.down_2(out)\n        out = F.max_pool2d(down2, kernel_size=2, stride=2)\n\n        down3 = self.down_3(out)\n        out = F.max_pool2d(down3, kernel_size=2, stride=2)\n\n        down4 = self.down_4(out)\n        out = F.max_pool2d(down4, kernel_size=2, stride=2)\n\n        down5 = self.down_5(out)\n        out = F.max_pool2d(down5, kernel_size=2, stride=2)\n\n        out = self.middle(out)\n        out = self.middle_t(out, output_size=down5.size())\n\n        out = torch.cat([down5, out], 1)\n        out = self.up_5(out)\n        out = self.up_5_t(out, output_size=down4.size())\n\n        out = torch.cat([down4, out], 1)\n        out = self.up_4(out)\n        out = self.up_4_t(out, output_size=down3.size())\n        \n        out = torch.cat([down3, out], 1)\n        out = self.up_3(out)\n        out = self.up_3_t(out, output_size=down2.size())\n        \n        out = torch.cat([down2, out], 1)\n        out = self.up_2(out)\n        out = self.up_2_t(out, output_size=down1.size())\n        \n        out = torch.cat([down1, out], 1)\n        out = self.up_1(out)\n        \n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Lost + Accuracy functions"},{"metadata":{"trusted":false},"cell_type":"code","source":"def iou_loss(y_true, y_pred):\n    '''\n    Intersection-Over-Union Loss\n    '''\n    y_true = torch.reshape(y_true, [-1])\n    y_pred = torch.reshape(y_pred, [-1])\n    intersection = (y_true * y_pred).sum()\n    score = (intersection + 1.) / (y_true.sum() + y_pred.sum() - intersection + 1.)\n    return 1 - score","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def iou_bce_loss(y_true, y_pred):\n    '''\n    Main loss function using:\n        Binary Cross Entropy +\n        Intersection-Over-Union Loss\n    '''\n    return 0.5 * F.binary_cross_entropy_with_logits(y_true, y_pred) + 0.5 * iou_loss(y_true, y_pred)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def mean_iou(y_true, y_pred, device):\n    '''\n    Mean-Intersection-Over-Union\n    '''\n    y_pred = torch.round(y_pred)\n    intersect = (y_true * y_pred).sum(axis=[1, 2, 3])\n    union = y_true.sum(axis=[1, 2, 3]) + y_pred.sum(axis=[1, 2, 3])\n    smooth = torch.ones(intersect.shape).to(device)\n    return ((intersect + smooth) / (union - intersect + smooth)).sum()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":false},"cell_type":"code","source":"net = Net().to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"optimizer = optim.Adam(net.parameters(), lr=0.0001)","execution_count":null,"outputs":[]},{"metadata":{"tags":["outputPrepend"],"trusted":false},"cell_type":"code","source":"for batch_ndx, sample in enumerate(train_gen):\n    x, y = sample\n\n    optimizer.zero_grad()\n\n    out = net(x)\n    loss = iou_bce_loss(out, y)\n    loss.backward()\n\n    optimizer.step()\n\n    iou = mean_iou(out, y, device)\n\n    print(\"Batch_ndx:{0:5d}, Loss:{1:2.4f}, Mean-Intersection-Over-Union:{2:2.4f}\"\n            .format(batch_ndx, loss, iou))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Conclusion"},{"metadata":{},"cell_type":"markdown","source":"Training this model from scratch will take hours and hours.<br>\nI suggest trying to use pre-trained custom ResNet50 and add your new layers + re-train with our dataset"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}