{"cells":[{"metadata":{"_uuid":"f16518bc621f64e59c33f11376ac8b1168b35be8"},"cell_type":"markdown","source":"# UNET Pytorch implementation\nThis notebook contains a custom UNET segmentation model that I implemented from scratch using pytorch, applied to the RSNA pneumonia challenge. The model was trained for 10 epochs (< 5 hours) on Google Cloud using 8 CPUs and 1 NVIDIA TESLA P100 GPU (specs: PyTorch 0.4.1, Python 3.6.3, CUDA 9.2.148.1, cuDNN 7.2.1).\nThe cnn architecture was inspired to [this model](https://github.com/ternaus/TernausNet), but adapted to a single-channel input and without using transfer learning. \n\nUnfortunately, it cannot be run on Kaggle using a batch_size of 25 images as in the original setup I ran on Google Cloud (GPU out of memory). \n**However, you can find the LB 0.123 submission file in the data attached to this kernel. I created a public dataset called \"pytorch-unet-pneumonia-output\" where I put the final submission file and some outputs obtained from the run on Google Cloud.**\n\nFeedback/questions are most welcome!\n"},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"%matplotlib inline\nimport os\nimport time\n\nimport skimage \nfrom skimage.transform import resize\nfrom skimage.exposure import rescale_intensity\nfrom scipy.ndimage.interpolation import map_coordinates\nfrom scipy.ndimage.filters import gaussian_filter\nimport PIL\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.dataset import Dataset as torchDataset\nimport torchvision as tv\nfrom torch.autograd import Variable\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib as mpl\nfrom matplotlib.patches import Rectangle\n\nimport shutil\n\nimport pydicom\n\nimport warnings\nwarnings.simplefilter(action='ignore', category=FutureWarning)\nwarnings.catch_warnings()\n\ngpu_available = True\n\noriginal_image_shape = 1024\n\ndatapath_orig = '../input/rsna-pneumonia-detection-challenge/'\ndatapath_prep = '../input/start-here-beginner-intro-to-lung-opacity-s1/'\ndatapath_out = '../input/pytorchunetpneumoniaoutput/'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"959d3096d11ae2c52c6a3072314fc9dc71fbb0ee"},"cell_type":"markdown","source":"Below I import the preprocessed labels data that I generated in a [previous eda kernel](https://www.kaggle.com/giuliasavorgnan/start-here-beginner-intro-to-lung-opacity-s1)."},{"metadata":{"_uuid":"d24af5228cf386fcadd225e442bd57afcfc9590b","trusted":true},"cell_type":"code","source":"# read train dataset\ndf_train = pd.read_csv(datapath_prep+'train.csv')\n# read test dataset\ndf_test = pd.read_csv(datapath_prep+'test.csv')\ndf_train.head(3)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8c82c8c09453c84b37718297d5d8feda2a5ed73a"},"cell_type":"markdown","source":"In an attempt to control the number of false positive boxes, I calculate the distribution of the boxes' area and manually select a lower limit for the unet model."},{"metadata":{"_uuid":"7a925348d5581f1ab749ba16c99e0273a9cce86f","trusted":true},"cell_type":"code","source":"# calculate minimum box area as benchmark for CNN model\ndf_train['box_area'] = df_train['width'] * df_train['height']\ndf_train['box_area'].hist(bins=100)\ndf_train['box_area'].describe()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"5692b292478bd2d534a8da94595717465181ef20","trusted":true},"cell_type":"code","source":"# arbitrary value for minimum box area in the CNN model\nmin_box_area = 10000","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"c7687d275d2641b47e05a373326e012b6892e92d"},"cell_type":"markdown","source":"The following code prepares the training data in a useful format for the unet model."},{"metadata":{"_uuid":"7fe5579f238cd9ddb2f51647128ff253fc5ccf8a","trusted":true},"cell_type":"code","source":"# shuffle and create patient ID list, then split into train and validation sets\nvalidation_frac = 0.10\n\ndf_train = df_train.sample(frac=1, random_state=42) # .sample(frac=1) does the shuffling\npIds = [pId for pId in df_train['patientId'].unique()]\n\npIds_valid = pIds[ : int(round(validation_frac*len(pIds)))]\npIds_train = pIds[int(round(validation_frac*len(pIds))) : ]\n\nprint('{} patient IDs shuffled and {}% of them used in validation set.'.format(len(pIds), validation_frac*100))\nprint('{} images went into train set and {} images went into validation set.'.format(len(pIds_train), len(pIds_valid)))\n\n# get test set patient IDs\npIds_test = df_test['patientId'].unique()\nprint('{} patient IDs in test set.'.format(len(pIds_test)))","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"0e44aae1c2ee31ebc8b836de4abeb5097f61ece4","trusted":true},"cell_type":"code","source":"def get_boxes_per_patient(df, pId):\n    '''\n    Given the dataset and one patient ID, \n    return an array of all the bounding boxes and their labels associated with that patient ID.\n    Example of return: \n    array([[x1, y1, width1, height1],\n           [x2, y2, width2, height2]])\n    '''\n    \n    boxes = df.loc[df['patientId']==pId][['x', 'y', 'width', 'height']].astype('int').values.tolist()\n    return boxes","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"ab8077059d5fe1142bc74171fbb2d7ba48dee932","trusted":true},"cell_type":"code","source":"# create dictionary of {patientId : list of boxes}\npId_boxes_dict = {}\nfor pId in df_train.loc[(df_train['Target']==1)]['patientId'].unique().tolist():\n    pId_boxes_dict[pId] = get_boxes_per_patient(df_train, pId)\nprint('{} ({:.1f}%) images have target boxes.'.format(len(pId_boxes_dict), 100*(len(pId_boxes_dict)/len(pIds))))","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"e961b8c68cfa33278acf5e39ad78b4511882f02a"},"cell_type":"markdown","source":"The following code pertains to the unet model proper."},{"metadata":{"_uuid":"1b49bb60eaa38c96f0d105bfc8d3a6aa5a698bda","trusted":true},"cell_type":"code","source":"# define a MinMaxScaler function for the images\ndef imgMinMaxScaler(img, scale_range):\n    \"\"\"\n    :param img: image to be rescaled\n    :param scale_range: (tuple) (min, max) of the desired rescaling\n    \"\"\"\n    warnings.filterwarnings(\"ignore\")\n    img = img.astype('float64')\n    img_std = (img - np.min(img)) / (np.max(img) - np.min(img))\n    img_scaled = img_std * float(scale_range[1] - scale_range[0]) + float(scale_range[0])\n    # round at closest integer and transform to integer \n    img_scaled = np.rint(img_scaled).astype('uint8')\n\n    return img_scaled","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"930a8b2aff636355c97a8cc2c773bb2bfa691763"},"cell_type":"code","source":"# define a \"warping\" image/mask function \ndef elastic_transform(image, alpha, sigma, random_state=None):\n    \"\"\"Elastic deformation of images as described in [Simard2003]_.\n    .. [Simard2003] Simard, Steinkraus and Platt, \"Best Practices for\n       Convolutional Neural Networks applied to Visual Document Analysis\", in\n       Proc. of the International Conference on Document Analysis and\n       Recognition, 2003.\n       Code adapted from https://gist.github.com/chsasank/4d8f68caf01f041a6453e67fb30f8f5a\n    \"\"\"\n    assert len(image.shape)==2, 'Image must have 2 dimensions.'\n\n    if random_state is None:\n        random_state = np.random.RandomState(None)\n\n    shape = image.shape\n\n    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode=\"constant\", cval=0) * alpha\n    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode=\"constant\", cval=0) * alpha\n\n    x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')\n    indices = np.reshape(x+dx, (-1, 1)), np.reshape(y+dy, (-1, 1))\n    \n    image_warped = map_coordinates(image, indices, order=1).reshape(shape)\n    \n    return image_warped","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"dc0cd3a5487265f5cd28b9abd9c0b329afd5b48b","trusted":true},"cell_type":"code","source":"# define the data generator class\nclass PneumoniaDataset(torchDataset):\n    \"\"\"\n        Pneumonia dataset that contains radiograph lung images as .dcm. \n        Each patient has one image named patientId.dcm.\n    \"\"\"\n\n    def __init__(self, root, subset, pIds, predict, boxes, rescale_factor=1, transform=None, rotation_angle=0, warping=False):\n        \"\"\"\n        :param root: it has to be a path to the folder that contains the dataset folders\n        :param subset: 'train' or 'test'\n        :param pIds: list of patient IDs\n        :param predict: boolean, if true returns images and target labels, otherwise returns only images\n        :param boxes: a {patientId : list of boxes} dictionary (ex: {'pId': [[x1, y1, w1, h1], [x2, y2, w2, h2]]}\n        :param rescale_factor: image rescale factor in network (image shape is supposed to be square)\n        :param transform: transformation applied to the images and their target masks\n        :param rotation_angle: float, defines range of random rotation angles for augmentation (-rotation_angle, +rotation_angle)\n        :param warping: boolean, whether applying augmentation warping to image\n        \"\"\"\n        \n        # initialize variables\n        self.root = os.path.expanduser(root)\n        self.subset = subset\n        if self.subset not in ['train', 'test']:\n            raise RuntimeError('Invalid subset ' + self.subset + ', it must be one of: \\'train\\' or \\'test\\'')\n        self.pIds = pIds\n        self.predict = predict\n        self.boxes = boxes\n        self.rescale_factor = rescale_factor\n        self.transform = transform\n        self.rotation_angle = rotation_angle\n        self.warping = warping\n\n        self.data_path = self.root + 'stage_1_'+self.subset+'_images/'\n        \n    def __getitem__(self, index):\n        # get the corresponding pId\n        pId = self.pIds[index]\n        # load dicom file as numpy array\n        img = pydicom.dcmread(os.path.join(self.data_path, pId+'.dcm')).pixel_array\n        # check if image is square\n        if (img.shape[0]!=img.shape[1]):\n            raise RuntimeError('Image shape {} should be square.'.format(img.shape))\n        original_image_shape = img.shape[0]\n        # calculate network image shape\n        image_shape = original_image_shape / self.rescale_factor\n        # check if image_shape is an integer\n        if (image_shape != int(image_shape)):\n            raise RuntimeError('Network image shape should be an integer.'.format(image_shape))\n        image_shape = int(image_shape)\n        # resize image \n        # IMPORTANT: skimage resize function rescales the output from 0 to 1, and pytorch doesn't like this!\n        # One solution would be using torchvision rescale function (but need to differentiate img and target transforms)\n        # Here I use skimage resize and then rescale the output again from 0 to 255\n        img = resize(img, (image_shape, image_shape), mode='reflect')\n        # recale image from 0 to 255\n        img = imgMinMaxScaler(img, (0,255))\n        # image warping augmentation\n        if self.warping:\n            img = elastic_transform(img, image_shape*2., image_shape*0.1)\n        # add trailing channel dimension\n        img = np.expand_dims(img, -1)\n        # apply rotation augmentation\n        if self.rotation_angle>0:\n            angle = self.rotation_angle * (2 * np.random.random_sample() - 1) # generate random angle \n            img = tv.transforms.functional.to_pil_image(img)\n            img = tv.transforms.functional.rotate(img, angle, resample=PIL.Image.BILINEAR)\n                                            \n        # apply transforms to image\n        if self.transform is not None:\n            img = self.transform(img)\n        \n        if not self.predict:\n            # create target mask\n            target = np.zeros((image_shape, image_shape))\n            # if patient ID has associated target boxes (=if image contains pneumonia)\n            if pId in self.boxes:\n                # loop through boxes\n                for box in self.boxes[pId]:\n                    # extract box coordinates \n                    x, y, w, h = box\n                    # rescale box coordinates\n                    x = int(round(x/rescale_factor))\n                    y = int(round(y/rescale_factor))\n                    w = int(round(w/rescale_factor))\n                    h = int(round(h/rescale_factor))\n                    # create a mask of 1s (255 is used because pytorch will rescale to 0-1) inside the box\n                    target[y:y+h, x:x+w] = 255 #\n                    target[target>255] = 255 # correct in case of overlapping boxes (shouldn't happen)\n            # add trailing channel dimension\n            target = np.expand_dims(target, -1)   \n            target = target.astype('uint8')\n            # apply rotation augmentation\n            if self.rotation_angle>0:\n                target = tv.transforms.functional.to_pil_image(target)\n                target = tv.transforms.functional.rotate(target, angle, resample=PIL.Image.BILINEAR)\n            # apply transforms to target\n            if self.transform is not None:\n                target = self.transform(target)\n            return img, target, pId\n        else: \n            return img, pId\n\n    def __len__(self):\n        return len(self.pIds)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"c6b0894fa817c5d72c1cef3af6e78658728d548c","trusted":true},"cell_type":"code","source":"# manual model parameters\nrescale_factor = 4 # resize factor to reduce image size (new_image_shape = original_image_shape / rescale_factor)\nbatch_size = 6 # I used 25 on GCP\n\n# recalculate minimum box area\nmin_box_area = int(round(min_box_area / float(rescale_factor**2)))\n\n# TBD add normalization of images into transforms\n# define transformation \ntransform = tv.transforms.Compose([tv.transforms.ToTensor()])\n\n# create datasets\ndataset_train = PneumoniaDataset(root=datapath_orig, subset='train', pIds=pIds_train, predict=False, \n                                 boxes=pId_boxes_dict, rescale_factor=rescale_factor, transform=transform,\n                                 rotation_angle=3, warping=True)\n\ndataset_valid = PneumoniaDataset(root=datapath_orig, subset='train', pIds=pIds_valid, predict=False, \n                                 boxes=pId_boxes_dict, rescale_factor=rescale_factor, transform=transform,\n                                 rotation_angle=0, warping=False)\n\ndataset_test = PneumoniaDataset(root=datapath_orig, subset='test', pIds=pIds_test, predict=True, \n                                boxes=None, rescale_factor=rescale_factor, transform=transform,\n                                rotation_angle=0, warping=False)\n\n# define the dataloaders with the previous dataset\nloader_train = DataLoader(dataset=dataset_train,\n                           batch_size=batch_size,\n                           shuffle=True) \n\nloader_valid = DataLoader(dataset=dataset_valid,\n                           batch_size=batch_size,\n                           shuffle=True) \n\nloader_test = DataLoader(dataset=dataset_test,\n                         batch_size=batch_size,\n                         shuffle=False) ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"2ca70aa468aaf893599aea5b4d566ea5b03f07e3","trusted":true},"cell_type":"code","source":"# Check if train images have been properly loaded\nprint('{} images in train set, {} images in validation set, and {} images in test set.'.format(len(dataset_train),\n                                                                                               len(dataset_valid),\n                                                                                               len(dataset_test)))\nimg_batch, target_batch, pId_batch = next(iter(loader_train))\nprint('Tensor batch size:', img_batch.size())\n\n# Display some examples\nfor i in np.random.choice(len(dataset_train), size=5, replace=False):\n    img, target, pId = dataset_train[i] # picking an image with pneumonia\n    print('\\nImage and mask shapes:', img.shape, target.shape)\n    print('Patient ID:', pId)\n    print('Image scale: {} - {}'.format(img[0].min(), img[0].max()))\n    print('Target mask scale: {} - {}'.format(target[0].min(), target[0].max()))\n    plt.imshow(img[0], cmap=mpl.cm.gist_gray) # [0] is the channel index (here there's just one channel)\n    plt.imshow(target[0], cmap=mpl.cm.jet, alpha=0.2)\n    plt.axis('off')\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"75273c4f0607b67877d79c00ad659b2bdb0fde93","trusted":true},"cell_type":"code","source":"# Check if test images have been properly loaded\nimg, pId = dataset_test[0] \nprint('Image shape:', img.shape)\nprint('Patient ID:', pId)\nprint('Image scale: {} - {}'.format(img[0].min(), img[0].max()))\nplt.imshow(img[0], cmap=mpl.cm.gist_gray) # [0] is the channel index (here there's just one channel)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"802485df954e99183280ab911a3f15abd184a3ef"},"cell_type":"markdown","source":"The basic block (conv_block) of the unet model consists of a [2D convolution - batch normalization - activation] sequence. The 2D convolution uses 3x3 filters with stride=1 and padding=1. The activation function is a leaky ReLU with alpha=0.03. \nNOTE: I have not attempted yet to optimize these hyperparameters.\n\nThe second function (conv_t_block) does the upsampling for the upscaling half of the unet."},{"metadata":{"_uuid":"18677b6eff4638b4a0186ec9e65237d48f572ddb","trusted":true},"cell_type":"code","source":"# Define the nn convolutional block\nclass 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# Define the nn transposed convolutional block\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":{"_uuid":"403551d53e6988d9f94534ce0461dec82ed86d1e","trusted":true},"cell_type":"code","source":"# the actual model function is defined here\n# NOTE: the comments are meant to help understand/check the input-output sizes of the tensor\n#       and they assume an input image size of 256x256, \n#       but the model can accept larger/smaller sizes (tho I haven't debugged it yet)\nclass PneumoniaUNET(nn.Module):\n\n    def __init__(self):\n        super(PneumoniaUNET, 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) # (1x256x256 -> 64x256x256)\n        out = F.max_pool2d(down1, kernel_size=2, stride=2) # (64x256x256 -> 64x128x128)\n\n        down2 = self.down_2(out) # (64x128x128 -> 128x128x128)\n        out = F.max_pool2d(down2, kernel_size=2, stride=2) # (128x128x128 -> 128x64x64)\n\n        down3 = self.down_3(out) # (128x64x64 -> 256x64x64)\n        out = F.max_pool2d(down3, kernel_size=2, stride=2) # (256x64x64 -> 256x32x32)\n\n        down4 = self.down_4(out) # (256x32x32 -> 512x32x32)\n        out = F.max_pool2d(down4, kernel_size=2, stride=2) # (512x32x32 -> 512x16x16)\n\n        down5 = self.down_5(out) # (512x16x16 -> 512x16x16)\n        out = F.max_pool2d(down5, kernel_size=2, stride=2) # (512x16x16 -> 512x8x8)\n\n        out = self.middle(out) # (512x8x8 -> 512x8x8)\n        out = self.middle_t(out, output_size=down5.size()) # (512x8x8 -> 256x16x16)\n\n        out = torch.cat([down5, out], 1) # (512x16x16-concat-256x16x16 -> 768x16x16)\n        out = self.up_5(out) # (768x16x16 -> 512x16x16)\n        out = self.up_5_t(out, output_size=down4.size()) # (512x16x16 -> 256x32x32)\n\n        out = torch.cat([down4, out], 1) # (512x32x32-concat-256x32x32 -> 768x32x32)\n        out = self.up_4(out) # (768x32x32 -> 512x32x32)\n        out = self.up_4_t(out, output_size=down3.size()) # (512x32x32 -> 128x64x64)\n        \n        out = torch.cat([down3, out], 1) # (256x64x64-concat-128x64x64 -> 384x64x64)\n        out = self.up_3(out) # (384x64x64 -> 256x64x64)\n        out = self.up_3_t(out, output_size=down2.size()) # (256x64x64 -> 64x128x128)\n        \n        out = torch.cat([down2, out], 1) # (128x128x128-concat-64x128x128 -> 192x128x128)\n        out = self.up_2(out) # (192x128x128 -> 128x128x128)\n        out = self.up_2_t(out, output_size=down1.size()) # (128x128x128 -> 32x256x256)\n        \n        out = torch.cat([down1, out], 1) # (64x256x256-concat-32x256x256 -> 96x256x256)\n        out = self.up_1(out) # (96x256x256 -> 1x256x256)\n        \n        return out    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"a40e724fbda4ad388f74aeee061b457d899977a0"},"cell_type":"code","source":"# # the actual model function is defined here\n# # NOTE: the comments are meant to help understand/check the input-output sizes of the tensor\n# #       and they assume an input image size of 256x256, \n# #       but the model can accept larger/smaller sizes (tho I haven't debugged it yet)\n# class PneumoniaUNET(nn.Module):\n\n#     def __init__(self):\n#         super(PneumoniaUNET, self).__init__()\n        \n#         self.down_1 = conv_block(in_channels=1, out_channels=64)\n#         self.down_2 = conv_block(in_channels=64, out_channels=128)\n#         self.down_3 = conv_block(in_channels=128, out_channels=256)\n#         self.down_4 = conv_block(in_channels=256, out_channels=512)\n#         self.down_5 = conv_block(in_channels=512, out_channels=512)\n\n#         self.middle = 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 = conv_block(in_channels=768, out_channels=512)\n#         self.up_5_t = conv_t_block(in_channels=512, out_channels=256)\n#         self.up_4 = conv_block(in_channels=768, out_channels=512)\n#         self.up_4_t = conv_t_block(in_channels=512, out_channels=128)\n#         self.up_3 = conv_block(in_channels=384, out_channels=256)\n#         self.up_3_t = conv_t_block(in_channels=256, out_channels=64)\n#         self.up_2 = conv_block(in_channels=192, out_channels=128)\n#         self.up_2_t = conv_t_block(in_channels=128, out_channels=32)\n#         self.up_1 = conv_block(in_channels=96, out_channels=1)\n        \n#     def forward(self, x):\n#         down1 = self.down_1(x) # (1x256x256 -> 64x256x256)\n#         out = F.max_pool2d(down1, kernel_size=2, stride=2) # (64x256x256 -> 64x128x128)\n\n#         down2 = self.down_2(out) # (64x128x128 -> 128x128x128)\n#         out = F.max_pool2d(down2, kernel_size=2, stride=2) # (128x128x128 -> 128x64x64)\n\n#         down3 = self.down_3(out) # (128x64x64 -> 256x64x64)\n#         out = F.max_pool2d(down3, kernel_size=2, stride=2) # (256x64x64 -> 256x32x32)\n\n#         down4 = self.down_4(out) # (256x32x32 -> 512x32x32)\n#         out = F.max_pool2d(down4, kernel_size=2, stride=2) # (512x32x32 -> 512x16x16)\n\n#         down5 = self.down_5(out) # (512x16x16 -> 512x16x16)\n#         out = F.max_pool2d(down5, kernel_size=2, stride=2) # (512x16x16 -> 512x8x8)\n\n#         out = self.middle(out) # (512x8x8 -> 512x8x8)\n#         out = self.middle_t(out, output_size=down5.size()) # (512x8x8 -> 256x16x16)\n\n#         out = torch.cat([down5, out], 1) # (512x16x16-concat-256x16x16 -> 768x16x16)\n#         out = self.up_5(out) # (768x16x16 -> 512x16x16)\n#         out = self.up_5_t(out, output_size=down4.size()) # (512x16x16 -> 256x32x32)\n\n#         out = torch.cat([down4, out], 1) # (512x32x32-concat-256x32x32 -> 768x32x32)\n#         out = self.up_4(out) # (768x32x32 -> 512x32x32)\n#         out = self.up_4_t(out, output_size=down3.size()) # (512x32x32 -> 128x64x64)\n        \n#         out = torch.cat([down3, out], 1) # (256x64x64-concat-128x64x64 -> 384x64x64)\n#         out = self.up_3(out) # (384x64x64 -> 256x64x64)\n#         out = self.up_3_t(out, output_size=down2.size()) # (256x64x64 -> 64x128x128)\n        \n#         out = torch.cat([down2, out], 1) # (128x128x128-concat-64x128x128 -> 192x128x128)\n#         out = self.up_2(out) # (192x128x128 -> 128x128x128)\n#         out = self.up_2_t(out, output_size=down1.size()) # (128x128x128 -> 32x256x256)\n        \n#         out = torch.cat([down1, out], 1) # (64x256x256-concat-32x256x256 -> 96x256x256)\n#         out = self.up_1(out) # (96x256x256 -> 1x256x256)\n        \n#         return out    ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"29d80e594d81b3cf1b4d31dd1dcd724e48757c20","trusted":true},"cell_type":"code","source":"# print model architecture\nprint(PneumoniaUNET())","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"c58d080e0979c01dd6edb4b7387e59dc04e759c4","trusted":true},"cell_type":"code","source":"# Create the loss function\n# Define the 2D Sigmoid + Binary Cross Entropy loss function BCEWithLogitsLoss\n# TBD add weights for unbalanced class\n# NOTE: Rather than using weights, I also intended to try implement focal loss (see RetinaNet)\nclass BCEWithLogitsLoss2d(nn.Module):\n\n    def __init__(self, weight=None, size_average=True):\n        super(BCEWithLogitsLoss2d, self).__init__()\n        self.loss = nn.BCEWithLogitsLoss(weight, size_average)\n\n    def forward(self, scores, targets):\n        scores_flat = scores.view(-1)\n        targets_flat = targets.view(-1)\n        return self.loss(scores_flat, targets_flat)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"73632678b3be947135bf03fa118643ebf8f7bf65","trusted":true},"cell_type":"code","source":"# Define auxiliary metric functions\n\n# define function that creates a square mask for a box from its coordinates \ndef box_mask(box, shape=1024):\n    \"\"\"\n    :param box: [x, y, w, h] box coordinates\n    :param shape: shape of the image (default set to maximum possible value, set to smaller to save memory)\n    :returns: (np.array of bool) mask as binary 2D array\n    \"\"\"\n    x, y, w, h = box\n    mask = np.zeros((shape, shape), dtype=bool)\n    mask[y:y+h, x:x+w] = True \n    return mask\n\n# # debug code for above function\n# plt.imshow(box_mask([5,20,50,100], shape=256), cmap=mpl.cm.jet)\n\n# define function that extracts confidence and coordinates of boxes from a prediction mask\ndef parse_boxes(msk, threshold=0.20, connectivity=None):\n    \"\"\"\n    :param msk: (torch.Tensor) CxWxH tensor representing the prediction mask\n    :param threshold: threshold in the range 0-1 above which a pixel is considered a positive target\n    :param connectivity: connectivity parameter for skimage.measure.label segmentation (can be None, 1, or 2)\n                         http://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label\n    :returns: (list, list) predicted_boxes, confidences\n    \"\"\"\n    # extract 2d array\n    msk = msk[0]\n    # select pixels above threshold and mark them as positives (1) in an array of equal size as the input prediction mask\n    pos = np.zeros(msk.shape)\n    pos[msk>threshold] = 1.\n    # label regions\n    lbl = skimage.measure.label(pos, connectivity=connectivity)\n    \n    predicted_boxes = []\n    confidences = []\n    # iterate over regions and extract box coordinates\n    for region in skimage.measure.regionprops(lbl):\n        # retrieve x, y, height and width and add to prediction string\n        y1, x1, y2, x2 = region.bbox\n        h = y2 - y1\n        w = x2 - x1\n        c = np.nanmean(msk[y1:y2, x1:x2])\n        # add control over box size (eliminate if too small)\n        if w*h > min_box_area: \n            predicted_boxes.append([x1, y1, w, h])\n            confidences.append(c)\n    \n    return predicted_boxes, confidences\n\n# # debug code for above function\n# plt.imshow(dataset_train[3][1][0], cmap=mpl.cm.jet) \n# print(dataset_train[3][1].shape)\n# print(parse_boxes(dataset_train[3][1]))\n\n# define function that creates prediction strings as expected in submission\ndef prediction_string(predicted_boxes, confidences):\n    \"\"\"\n    :param predicted_boxes: [[x1, y1, w1, h1], [x2, y2, w2, h2], ...] list of predicted boxes coordinates \n    :param confidences: [c1, c2, ...] list of confidence values for the predicted boxes\n    :returns: prediction string 'c1 x1 y1 w1 h1 c2 x2 y2 w2 h2 ...'\n    \"\"\"\n    prediction_string = ''\n    for c, box in zip(confidences, predicted_boxes):\n        prediction_string += ' ' + str(c) + ' ' + ' '.join([str(b) for b in box])\n    return prediction_string[1:]   \n\n# # debug code for above function\n# predicted_boxes, confidences = parse_boxes(dataset_train[3][1])\n# print(predicted_boxes, confidences)\n# print(prediction_string(predicted_boxes, confidences))\n\n# define iou function\ndef IoU(pr, gt):\n    \"\"\"\n    :param pr: (numpy_array(bool)) prediction array \n    :param gt: (numpy_array(bool)) ground truth array \n    :returns: IoU (pr, gt) = intersection (pr, gt) / union (pr, gt)\n    \"\"\"\n    IoU = (pr & gt).sum() / ((pr | gt).sum() + 1.e-9)\n    return IoU\n\n# # debug code for above function\n# pr = box_mask([50,60,100,150], shape=256)\n# gt = box_mask([30,40,100,140], shape=256)\n# plt.imshow(pr, cmap=mpl.cm.Reds, alpha=0.3)\n# plt.imshow(gt, cmap=mpl.cm.Greens, alpha=0.3)\n# print(IoU(pr, gt))\n\n# define precision function\ndef precision(tp, fp, fn):\n    \"\"\"\n    :param tp: (int) number of true positives\n    :param fp: (int) number of false positives\n    :param fn: (int) number of false negatives\n    :returns: precision metric for one image at one threshold\n    \"\"\"\n    return float(tp) / (tp + fp + fn + 1.e-9)\n\n# # debug code for above function\n# print(precision(3,1,1))\n\n# define function that calculates the average precision of an image\ndef average_precision_image(predicted_boxes, confidences, target_boxes, shape=1024):\n    \"\"\"\n    :param predicted_boxes: [[x1, y1, w1, h1], [x2, y2, w2, h2], ...] list of predicted boxes coordinates \n    :param confidences: [c1, c2, ...] list of confidence values for the predicted boxes\n    :param target_boxes: [[x1, y1, w1, h1], [x2, y2, w2, h2], ...] list of target boxes coordinates \n    :param shape: shape of the boolean masks (default set to maximum possible value, set to smaller to save memory)\n    :returns: average_precision\n    \"\"\"\n    \n    # if both predicted and target boxes are empty, precision is NaN (and doesn't count towards the batch average)\n    if predicted_boxes == [] and target_boxes == []:\n        return np.nan\n    else:\n        # if we have predicted boxes but no target boxes, precision is 0\n        if len(predicted_boxes)>0 and target_boxes == []:\n            return 0.0\n        # if we have target boxes but no predicted boxes, precision is 0\n        elif len(target_boxes)>0 and predicted_boxes == []:\n            return 0.0\n        # if we have both predicted and target boxes, proceed to calculate image average precision\n        else:\n            # define list of thresholds for IoU [0.4 , 0.45, 0.5 , 0.55, 0.6 , 0.65, 0.7 , 0.75]\n            thresholds = np.arange(0.4, 0.8, 0.05) \n            # sort boxes according to their confidence (from largest to smallest)\n            predicted_boxes_sorted = list(reversed([b for _, b in sorted(zip(confidences, predicted_boxes), \n                                                                         key=lambda pair: pair[0])]))            \n            average_precision = 0.0\n            for t in thresholds: # iterate over thresholds\n                # with a first loop we measure true and false positives\n                tp = 0 # initiate number of true positives\n                fp = len(predicted_boxes) # initiate number of false positives \n                for box_p in predicted_boxes_sorted: # iterate over predicted boxes coordinates\n                    box_p_msk = box_mask(box_p, shape) # generate boolean mask\n                    for box_t in target_boxes: # iterate over ground truth boxes coordinates\n                        box_t_msk = box_mask(box_t, shape) # generate boolean mask\n                        iou = IoU(box_p_msk, box_t_msk) # calculate IoU\n                        if iou>t:\n                            tp += 1 # if IoU is above the threshold, we got one more true positive\n                            fp -= 1 # and one less false positive\n                            break # proceed to the next predicted box\n                # with a second loop we measure false negatives\n                fn = len(target_boxes) # initiate number of false negatives\n                for box_t in target_boxes: # iterate over ground truth boxes coordinates\n                    box_t_msk = box_mask(box_t, shape) # generate boolean mask\n                    for box_p in predicted_boxes_sorted: # iterate over predicted boxes coordinates\n                        box_p_msk = box_mask(box_p, shape) # generate boolean mask\n                        iou = IoU(box_p_msk, box_t_msk) # calculate IoU\n                        if iou>t:\n                            fn -= 1\n                            break # proceed to the next ground truth box\n                # TBD: this algo must be checked against the official Kaggle evaluation method which is still not clear...\n                average_precision += precision(tp, fp, fn) / float(len(thresholds))\n            return average_precision\n\n# # debug code for above function\n# confidences = [0.3, 0.9]\n# predicted_boxes = [[20,20,60,70], [110,110,50,70]]\n# target_boxes = [[25,25,60,70], [100,100,50,70]]#, [200, 200, 30, 50]]\n# for box_p in predicted_boxes:\n#     plt.imshow(box_mask(box_p, shape=256), cmap=mpl.cm.Reds, alpha=0.3)\n# for box_t in target_boxes:\n#     plt.imshow(box_mask(box_t, shape=256), cmap=mpl.cm.Greens, alpha=0.3)\n# print(average_precision_image(predicted_boxes, confidences, target_boxes)) \n\n# define function that calculates the average precision of a batch of images\ndef average_precision_batch(output_batch, pIds, pId_boxes_dict, rescale_factor, shape=1024, return_array=False):\n    \"\"\"\n    :param output_batch: cnn model output batch\n    :param pIds: (list) list of patient IDs contained in the output batch\n    :param rescale_factor: CNN image rescale factor\n    :param shape: shape of the boolean masks (default set to maximum possible value, set to smaller to save memory)\n    :returns: average_precision\n    \"\"\"\n    \n    batch_precisions = []\n    for msk, pId in zip(output_batch, pIds): # iterate over batch prediction masks and relative patient IDs\n        # retrieve target boxes from dictionary (quicker than from mask itself)\n        target_boxes = pId_boxes_dict[pId] if pId in pId_boxes_dict else []\n        # rescale coordinates of target boxes\n        if len(target_boxes)>0:\n            target_boxes = [[int(round(c/float(rescale_factor))) for c in box_t] for box_t in target_boxes]\n        # extract prediction boxes and confidences\n        predicted_boxes, confidences = parse_boxes(msk) \n        batch_precisions.append(average_precision_image(predicted_boxes, confidences, target_boxes, shape=shape))\n    if return_array:\n        return np.asarray(batch_precisions)\n    else:\n        return np.nanmean(np.asarray(batch_precisions)) \n\n# # debug code for above function\n# targets = []\n# pIds = []\n# for i in range(5):\n#     (img, target, pId) = dataset_train[i]\n#     targets.append(target)\n#     pIds.append(pId)\n# # targets[0] = targets[1] #or pIds[0] = 'nan'\n# average_precision_batch(targets, pIds, pId_boxes_dict, rescale_factor, shape=256)    \n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"207499baf520b5952e347ff28351d66b89f1d7ca","trusted":true},"cell_type":"code","source":"class RunningAverage():\n    \"\"\"A simple class that maintains the running average of a quantity\n    \n    Example:\n    ```\n    loss_avg = RunningAverage()\n    loss_avg.update(2)\n    loss_avg.update(4)\n    loss_avg() = 3\n    ```\n    \"\"\"\n    def __init__(self):\n        self.steps = 0\n        self.total = 0\n    \n    def update(self, val):\n        self.total += val\n        self.steps += 1\n    \n    def __call__(self):\n        return self.total/float(self.steps)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"4c14e5666c2ddbb03edcc139bce0a2c445ee8e21","trusted":true},"cell_type":"code","source":"def save_checkpoint(state, is_best, metric):\n    \"\"\"Saves model and training parameters at 'last.pth.tar'. If is_best==True, also saves\n    'best.pth.tar'\n    Args:\n        state: (dict) contains model's state_dict, may contain other keys such as epoch, optimizer state_dict\n        is_best: (bool) True if it is the best model seen till now\n    \"\"\"\n    filename = 'last.pth.tar'\n    torch.save(state, filename)\n    if is_best:\n        shutil.copyfile(filename, metric+'.best.pth.tar')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d9346c5e323d2ce3281044b8049abfb0d8234c49","trusted":true},"cell_type":"code","source":"# define the training function\ndef train(model, dataloader, optimizer, loss_fn, num_steps, pId_boxes_dict, rescale_factor, shape, save_summary_steps=5):\n    # set model to train model\n    model.train()\n\n    summary = []\n    loss_avg = RunningAverage()\n\n    loss_avg_t_hist_ep, loss_t_hist_ep, prec_t_hist_ep = [], [], []\n\n    # iterate over batches\n    start = time.time()        \n    \n    for i, (input_batch, labels_batch, pIds_batch) in enumerate(dataloader):\n        # break loop after num_steps batches (useful for debugging)\n        if i > num_steps:\n            break\n        # Convert torch tensor to Variable\n        input_batch = Variable(input_batch).cuda(async=True) if gpu_available else Variable(input_batch).float()\n        labels_batch = Variable(labels_batch).cuda(async=True) if gpu_available else Variable(labels_batch).float()\n            \n        # compute output\n        optimizer.zero_grad()\n        output_batch = model(input_batch)\n\n        # compute loss\n        loss = loss_fn(output_batch, labels_batch)\n\n        # compute gradient and do optimizer step\n        loss.backward()\n        optimizer.step()\n\n        # update loss running average\n        loss_avg.update(loss.item())\n        loss_t_hist_ep.append(loss.item())\n        loss_avg_t_hist_ep.append(loss_avg())\n\n        # Evaluate summaries only once in a while\n        if i % save_summary_steps == 0:\n            # extract data from torch Variable, move to cpu\n            output_batch = output_batch.data.cpu().numpy()\n            # compute average precision on this batch\n            prec_batch = average_precision_batch(output_batch, pIds_batch, pId_boxes_dict, rescale_factor, shape)\n            prec_t_hist_ep.append(prec_batch)\n            # log results\n            summary_batch_string = \"batch loss = {:05.7f} ;  \".format(loss.item())\n            summary_batch_string += \"average loss = {:05.7f} ;  \".format(loss_avg())\n            summary_batch_string += \"batch precision = {:05.7f} ;  \".format(prec_batch)\n            print('--- Train batch {} / {}: '.format(i, num_steps) + summary_batch_string)\n            delta_time = time.time() - start\n            print('    {} batches processed in {:.2f} seconds'.format(save_summary_steps, delta_time))\n            start = time.time()\n\n    # log epoch summary\n    metrics_string = \"average loss = {:05.7f} ;  \".format(loss_avg())\n    print(\"- Train epoch metrics summary: \" + metrics_string)\n    \n    return loss_avg_t_hist_ep, loss_t_hist_ep, prec_t_hist_ep","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"9de116575eb50215b6288249233d4145fa27c388","trusted":true},"cell_type":"code","source":"def evaluate(model, dataloader, loss_fn, num_steps, pId_boxes_dict, rescale_factor, shape):\n\n    # set model to evaluation mode\n    model.eval()\n\n    losses = []\n    precisions = []\n\n    # compute metrics over the dataset\n    start = time.time()\n    for i, (input_batch, labels_batch, pIds_batch) in enumerate(dataloader):\n        # break loop after num_steps batches (useful for debugging)\n        if i > num_steps:\n            break\n        # Convert torch tensor to Variable\n        input_batch = Variable(input_batch).cuda(async=True) if gpu_available else Variable(input_batch).float()\n        labels_batch = Variable(labels_batch).cuda(async=True) if gpu_available else Variable(labels_batch).float()\n\n        # compute model output\n        output_batch = model(input_batch)\n        # compute loss of batch\n        loss = loss_fn(output_batch, labels_batch)\n        losses.append(loss.item())\n\n        # extract data from torch Variable, move to cpu\n        output_batch = output_batch.data.cpu()\n        # compute individual precisions of batch images\n        prec_batch = average_precision_batch(output_batch, pIds_batch, pId_boxes_dict, rescale_factor, shape, return_array=True)\n        for p in prec_batch:\n            precisions.append(p)\n        print('--- Validation batch {} / {}: '.format(i, num_steps))\n\n    # compute mean of all metrics in summary\n    metrics_mean = {'loss' : np.nanmean(losses),\n                    'precision' : np.nanmean(np.asarray(precisions))}\n    metrics_string = \"average loss = {:05.7f} ;  \".format(metrics_mean['loss'])\n    metrics_string += \"average precision = {:05.7f} ;  \".format(metrics_mean['precision'])\n    print(\"- Eval metrics : \" + metrics_string)\n    delta_time = time.time() - start\n    print('  Evaluation run in {:.2f} seconds.'.format(delta_time))\n    \n    return metrics_mean","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"2e5cff316bb314cbb22b977864dedb664f99d829","trusted":true},"cell_type":"code","source":"def train_and_evaluate(model, train_dataloader, val_dataloader, lr_init, loss_fn, num_epochs, \n                       num_steps_train, num_steps_eval, pId_boxes_dict, rescale_factor, shape, restore_file=None):\n\n    # reload weights from restore_file if specified\n    if restore_file is not None:\n        checkpoint = torch.load(restore_file)\n        model.load_state_dict(checkpoint['state_dict'])\n            \n    best_val_loss = 1e+15\n    best_val_prec = 0.0\n    best_loss_model = None\n    best_prec_model = None\n\n    loss_t_history = []\n    loss_v_history = []\n    loss_avg_t_history = []\n    prec_t_history = []\n    prec_v_history = []\n\n    for epoch in range(num_epochs):\n        start = time.time()\n        \n        # define the optimizer\n        lr = lr_init * 0.5**float(epoch) # reduce the learning rate at each epoch\n        optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n        # Run one epoch\n        print(\"Epoch {}/{}. Learning rate = {:05.3f}.\".format(epoch + 1, num_epochs, lr))\n\n        # train model for a whole epoc (one full pass over the training set)\n        loss_avg_t_hist_ep, loss_t_hist_ep, prec_t_hist_ep = train(model, train_dataloader, optimizer, loss_fn, \n                                                                   num_steps_train, pId_boxes_dict, rescale_factor, shape)\n        loss_avg_t_history += loss_avg_t_hist_ep\n        loss_t_history += loss_t_hist_ep\n        prec_t_history += prec_t_hist_ep\n        \n        # Evaluate for one epoch on validation set\n        val_metrics = evaluate(model, val_dataloader, loss_fn, num_steps_eval, pId_boxes_dict, rescale_factor, shape)\n\n        val_loss = val_metrics['loss']\n        val_prec = val_metrics['precision']\n        \n        loss_v_history += len(loss_t_hist_ep) * [val_loss]\n        prec_v_history += len(prec_t_hist_ep) * [val_prec]\n\n        is_best_loss = val_loss<=best_val_loss\n        is_best_prec = val_prec>=best_val_prec\n        \n        if is_best_loss:\n            print(\"- Found new best loss: {:.4f}\".format(val_loss))\n            best_val_loss = val_loss\n            best_loss_model = model\n        if is_best_prec:\n            print(\"- Found new best precision: {:.4f}\".format(val_prec))\n            best_val_prec = val_prec\n            best_prec_model = model\n            \n        # Save best weights based on best_val_loss and best_val_prec\n        save_checkpoint({'epoch': epoch + 1,\n                         'state_dict': model.state_dict(),\n                         'optim_dict' : optimizer.state_dict()},\n                         is_best=is_best_loss,\n                         metric='loss')\n        save_checkpoint({'epoch': epoch + 1,\n                         'state_dict': model.state_dict(),\n                         'optim_dict' : optimizer.state_dict()},\n                         is_best=is_best_prec,\n                         metric='prec')\n        \n        delta_time = time.time() - start\n        print('Epoch run in {:.2f} minutes'.format(delta_time/60.))\n\n    histories = {'loss avg train' : loss_avg_t_history,\n                 'loss train' : loss_t_history,\n                 'precision train' : prec_t_history,\n                 'loss validation' : loss_v_history, \n                 'precision validation' : prec_v_history}\n    best_models = {'best loss model' : best_loss_model,\n                   'best precision model' : best_prec_model}\n    \n    return histories, best_models","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"5c1e0f2d944a30cfa72d39b118caa72eba6286a5","trusted":true},"cell_type":"code","source":"def predict(model, dataloader): \n\n    # set model to evaluation mode\n    model.eval()\n    \n    predictions = {}\n\n    for i, (test_batch, pIds) in enumerate(dataloader):\n        print('Predicting batch {} / {}.'.format(i+1, len(dataloader)))\n        # Convert torch tensor to Variable\n        test_batch = Variable(test_batch).cuda(async=True) if gpu_available else Variable(test_batch).float()\n            \n        # compute output\n        output_batch = model(test_batch)\n        sig = nn.Sigmoid().cuda()\n        output_batch = sig(output_batch)\n        output_batch = output_batch.data.cpu().numpy()\n        for pId, output in zip(pIds, output_batch):\n            predictions[pId] = output\n        \n    return predictions","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"e0a7bcf8f26561d276caa4ee7687e04c4ec0461a","trusted":true},"cell_type":"code","source":"# train and evaluate the model\ndebug = False\n\n# define an instance of the model\nmodel = PneumoniaUNET().cuda() if gpu_available else PneumoniaUNET()\n# define the loss function\nloss_fn = BCEWithLogitsLoss2d().cuda() if gpu_available else BCEWithLogitsLoss2d()\n# define initial learning rate (will be reduced over epochs)\nlr_init = 0.5\n\nnum_epochs = 2 if debug else 2\nnum_steps_train = 50 if debug else len(loader_train)\nnum_steps_eval = 10 if debug else len(loader_valid)\n\nshape = int(round(original_image_shape / rescale_factor))\n\n# Train the model\nprint(\"Starting training for {} epochs\".format(num_epochs))\nhistories, best_models = train_and_evaluate(model, loader_train, loader_valid, lr_init, loss_fn, \n                                            num_epochs, num_steps_train, num_steps_eval, pId_boxes_dict, rescale_factor, shape)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"27b584a443adff0f3c8a47792bba5c5756cd3b4b","trusted":true},"cell_type":"code","source":"# visualize training loss history\nplt.plot(range(len(histories['loss train'])), histories['loss train'], color='k', label='loss train')\nplt.plot(range(len(histories['loss avg train'])), histories['loss avg train'], color='g', ls='dashed', label='loss avg train')\nplt.plot(range(len(histories['loss validation'])), histories['loss validation'], color='r', label='loss validation')\nplt.legend()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"be45f54706df3e60ee3e652f73758c214bf48e4a","trusted":true},"cell_type":"code","source":"plt.plot(range(len(histories['precision train'])), histories['precision train'], color='k', label='precision train')\nplt.plot(range(len(histories['precision validation'])), histories['precision validation'], color='r', label='precision validation')\nplt.legend()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"7800bc360cd836bd95464d45214ddfc245c37f38","trusted":true},"cell_type":"code","source":"# pick model with best precision\nbest_model = best_models['best precision model']","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"e43d8eaf21f7ff546e2102e111f9d4586a1464a2","trusted":true},"cell_type":"code","source":"# create predictions for the validation set to compare against ground truth\ndataset_valid = PneumoniaDataset(root=datapath_orig, subset='train', pIds=pIds_valid, predict=True, \n                                 boxes=None, rescale_factor=rescale_factor, transform=transform)\nloader_valid = DataLoader(dataset=dataset_valid,\n                          batch_size=batch_size,\n                          shuffle=False) \n\npredictions_valid = predict(best_model, loader_valid)\nprint('Predicted {} validation images.'.format(len(predictions_valid)))","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"71801746920ea5effb8d7d94f5921b97569b45f7","trusted":true},"cell_type":"code","source":"def rescale_box_coordinates(box, rescale_factor):\n    x, y, w, h = box\n    x = int(round(x/rescale_factor))\n    y = int(round(y/rescale_factor))\n    w = int(round(w/rescale_factor))\n    h = int(round(h/rescale_factor))\n    return [x, y, w, h]    ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"54b741a09c36add0d71e73481fcca7ca0c06b820","trusted":true},"cell_type":"code","source":"def draw_boxes(predicted_boxes, confidences, target_boxes, ax, angle=0):\n    if len(predicted_boxes)>0:\n        for box, c in zip(predicted_boxes, confidences):\n            # extracting individual coordinates\n            x, y, w, h = box \n            # create a rectangle patch\n            patch = Rectangle((x,y), w, h, color='red', ls='dashed',\n                              angle=angle, fill=False, lw=4, joinstyle='round', alpha=0.6)\n            # get current axis and draw rectangle\n            ax.add_patch(patch)\n            # add confidence value in annotation text\n            ax.text(x+w/2., y-5, '{:.2}'.format(c), color='red', size=20, va='center', ha='center')\n    if len(target_boxes)>0:\n        for box in target_boxes:\n            # rescale and extract box coordinates\n            x, y, w, h = box\n            # create a rectangle patch\n            patch = Rectangle((x,y), w, h, color='red',  \n                              angle=angle, fill=False, lw=4, joinstyle='round', alpha=0.6)\n            # get current axis and draw rectangle\n            ax.add_patch(patch)\n    \n    return ax","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"b1df63cc6f3fc4d62db21006d7e50693dfe0d67c","trusted":false},"cell_type":"code","source":"# grid search to cross-validate the best threshold for the boxes \nbest_threshold = None\nbest_avg_precision_valid = 0.0\nthresholds = np.arange(0.01, 0.60, 0.01)\navg_precision_valids = []\nfor threshold in thresholds:\n    precision_valid = []\n    for i in range(len(dataset_valid)):\n        img, pId = dataset_valid[i]\n        target_boxes = [rescale_box_coordinates(box, rescale_factor) for box in pId_boxes_dict[pId]] if pId in pId_boxes_dict else []\n        prediction = predictions_valid[pId]\n        predicted_boxes, confidences = parse_boxes(prediction, threshold=threshold, connectivity=None)\n        avg_precision_img = average_precision_image(predicted_boxes, confidences, target_boxes, shape=img[0].shape[0])\n        precision_valid.append(avg_precision_img)\n    avg_precision_valid = np.nanmean(precision_valid)\n    avg_precision_valids.append(avg_precision_valid)\n    print('Threshold: {}, average precision validation: {:03.5f}'.format(threshold, avg_precision_valid))\n    if avg_precision_valid>best_avg_precision_valid:\n        print('Found new best average precision validation!')\n        best_avg_precision_valid = avg_precision_valid\n        best_threshold = threshold\nplt.plot(thresholds, avg_precision_valids)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"6e06f1b38cc25b40f3b628fd770e182dd31097e2"},"cell_type":"code","source":"# check the results on the validation set\nfor i in range(len(dataset_valid)):\n    img, pId = dataset_valid[i]\n    target_boxes = [rescale_box_coordinates(box, rescale_factor) for box in pId_boxes_dict[pId]] if pId in pId_boxes_dict else []\n    prediction = predictions_valid[pId]\n    predicted_boxes, confidences = parse_boxes(prediction, threshold=best_threshold, connectivity=None)\n    avg_precision_img = average_precision_image(predicted_boxes, confidences, target_boxes, shape=img[0].shape[0])\n    if i%100==0: # print every 100\n        plt.imshow(img[0], cmap=mpl.cm.gist_gray) # [0] is the channel index (here there's just one channel)\n        plt.imshow(prediction[0], cmap=mpl.cm.jet, alpha=0.5)\n        draw_boxes(predicted_boxes, confidences, target_boxes, plt.gca())\n        print('Prediction mask scale:', prediction[0].min(), '-', prediction[0].max())\n        print('Prediction string:', prediction_string(predicted_boxes, confidences))\n        print('Ground truth boxes:', target_boxes)\n        print('Average precision image: {:05.5f}'.format(avg_precision_img))\n        plt.show()\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"116af95be7c76068865aebcb1833fc65682ea583","trusted":true},"cell_type":"code","source":"# create submission predictions for the test set\npredictions_test = predict(best_model, loader_test)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"a6630b06e78a0646ca232c6b096a416f4a8f1e27","trusted":true},"cell_type":"code","source":"print('Predicted {} images.'.format(len(predictions_test)))\nfor k, v in predictions_test.items():\n    print(v.shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"7f6e6cffb7aefc82af074d0706d5a4b3b6c92bf8","trusted":true},"cell_type":"code","source":"df_sub = df_test[['patientId']].copy(deep=True)\ndef get_prediction_string_per_pId(pId):\n    prediction = predictions_test[pId]\n    predicted_boxes, confidences = parse_boxes(prediction, threshold=best_threshold, connectivity=None)\n    predicted_boxes = [rescale_box_coordinates(box, 1/rescale_factor) for box in predicted_boxes]\n    return prediction_string(predicted_boxes, confidences)\ndf_sub['predictionString'] = df_sub['patientId'].apply(lambda x: get_prediction_string_per_pId(x) if x in pIds_test else '')\nprint('Number of non null prediction strings: {} ({:05.2f}%)'.format(df_sub.loc[df_sub['predictionString']!=''].shape[0],\n                                                    100. * df_sub.loc[df_sub['predictionString']!=''].shape[0]/df_sub.shape[0]))\ndf_sub.head(10)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"c83de59cf3d57ab80fb5a8d23f884bf9706cf961","trusted":true},"cell_type":"code","source":"df_sub.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"865f866c851c9a0917d86813eddbe9aafcbd23e5","trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}