{"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":"code","source":"# Libraries\nimport os\nfrom os.path import join\nfrom tqdm import tqdm\nimport random\n\nimport numpy as np\nimport pandas as pd\npd.set_option('display.max_rows', 500)\npd.set_option('display.max_columns', 500)\npd.set_option('display.width', 1000)\nimport matplotlib.pyplot as plt\nplt.rcParams.update({'font.size': 18})\nimport cv2\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset, sampler\n\nfrom albumentations import (HorizontalFlip, VerticalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise)\n\ndef initialize_seeds(seed):\n    np.random.seed(seed)\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    \ninitialize_seeds(2021)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-11-04T20:31:14.925655Z","iopub.execute_input":"2021-11-04T20:31:14.926329Z","iopub.status.idle":"2021-11-04T20:31:18.56768Z","shell.execute_reply.started":"2021-11-04T20:31:14.926147Z","shell.execute_reply":"2021-11-04T20:31:18.566695Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Dataset\n\n### Understand the Structure of the Dataset\n\n\n*   train - train images in PNG format\n*   The training annotations -> run length encoded masks\n*   Images -> PNG format (The number of images is small, but the number of annotated objects is quite high.)\n*   Test set -> 240 images\n\n### Files\n\n**train.csv** - IDs and masks for all training objects. None of this metadata is provided for the test set.\n* id - unique identifier for object\n* annotation - run length encoded pixels for the identified neuronal cell\n* width - source image width\n* height - source image height\n* cell_type - the cell line\n* plate_time - time plate was created\n* sample_date - date sample was created\n* sample_id - sample identifier\n* elapsed_timedelta - time since first image taken of sample\n\n**sample_submission.csv** - a sample submission file in the correct format\n\n**train** - train images in PNG format\n\n**test** - test images in PNG format. Only a few test set images are available for download; the remainder can only be accessed by your notebooks when you submit.\n\n**train_semi_supervised** - unlabeled images offered in case you want to use additional data for a semi-supervised approach.\n\n**LIVECell_dataset_2021** - A mirror of the data from the LIVECell dataset. LIVECell is the predecessor dataset to this competition. You will find extra data for the SH-SHY5Y cell line, plus several other cell lines not covered in the competition dataset that may be of interest for transfer learning.","metadata":{}},{"cell_type":"code","source":"DATA_PATH             = '../input/sartorius-cell-instance-segmentation'\nSAMPLE_SUBMISSION     = join(DATA_PATH,'train')\nTRAIN_CSV             = join(DATA_PATH,'train.csv')\nTRAIN_PATH            = join(DATA_PATH,'train')\nTEST_PATH             = join(DATA_PATH,'test')\n\ndf_train = pd.read_csv(TRAIN_CSV)\nprint(f'Training Set Shape: {df_train.shape} - {df_train[\"id\"].nunique()} \\\nImages - Memory Usage: {df_train.memory_usage().sum() / 1024 ** 2:.2f} MB')","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:18.569603Z","iopub.execute_input":"2021-11-04T20:31:18.56988Z","iopub.status.idle":"2021-11-04T20:31:19.250465Z","shell.execute_reply.started":"2021-11-04T20:31:18.569848Z","shell.execute_reply":"2021-11-04T20:31:19.248003Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_decode(mask_rle, shape, color=1):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0] * shape[1], dtype=np.float32)\n    for lo, hi in zip(starts, ends):\n        img[lo : hi] = color\n    return img.reshape(shape)\n\ndef build_masks(df_train, image_id, input_shape):\n    height, width = input_shape\n    labels = df_train[df_train[\"id\"] == image_id][\"annotation\"].tolist()\n    mask = np.zeros((height, width))\n    for label in labels:\n        mask += rle_decode(label, shape=(height, width))\n    mask = mask.clip(0, 1)\n    return np.array(mask)","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:19.252531Z","iopub.execute_input":"2021-11-04T20:31:19.252891Z","iopub.status.idle":"2021-11-04T20:31:19.266077Z","shell.execute_reply.started":"2021-11-04T20:31:19.252845Z","shell.execute_reply":"2021-11-04T20:31:19.265105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nclass CellDataset(Dataset):\n    def __init__(self, df: pd.core.frame.DataFrame, train:bool):\n        self.IMAGE_RESIZE = (224, 224)\n        self.RESNET_MEAN = (0.485, 0.456, 0.406)\n        self.RESNET_STD = (0.229, 0.224, 0.225)\n        self.df = df\n        self.base_path = TRAIN_PATH\n        self.gb = self.df.groupby('id')\n        self.transforms = Compose([Resize( self.IMAGE_RESIZE[0],  self.IMAGE_RESIZE[1]), \n                                   Normalize(mean=self.RESNET_MEAN, std= self.RESNET_STD, p=1), \n                                   HorizontalFlip(p=0.5),\n                                   VerticalFlip(p=0.5)])\n        \n        # Split train and val set\n        all_image_ids = np.array(df_train.id.unique())\n        np.random.seed(42)\n        iperm = np.random.permutation(len(all_image_ids))\n        num_train_samples = int(len(all_image_ids) * 0.9)\n\n        if train:\n            self.image_ids = all_image_ids[iperm[:num_train_samples]]\n        else:\n             self.image_ids = all_image_ids[iperm[num_train_samples:]]\n\n    def __getitem__(self, idx: int) -> dict:\n\n        image_id = self.image_ids[idx]\n        df = self.gb.get_group(image_id)\n\n        # Read image\n        image_path = os.path.join(self.base_path, image_id + \".png\")\n        image = cv2.imread(image_path)\n\n        # Create the mask\n        mask = build_masks(df_train, image_id, input_shape=(520, 704))\n        mask = (mask >= 1).astype('float32')\n        augmented = self.transforms(image=image, mask=mask)\n        image = augmented['image']\n        mask = augmented['mask']\n        # print(np.moveaxis(image,0,2).shape)\n        return np.moveaxis(np.array(image),2,0), mask.reshape((1, self.IMAGE_RESIZE[0], self.IMAGE_RESIZE[1]))\n\n\n    def __len__(self):\n        return len(self.image_ids)","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:19.26847Z","iopub.execute_input":"2021-11-04T20:31:19.268944Z","iopub.status.idle":"2021-11-04T20:31:19.288831Z","shell.execute_reply.started":"2021-11-04T20:31:19.268902Z","shell.execute_reply":"2021-11-04T20:31:19.287962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_train = CellDataset(df_train, train=True)\ndl_train = DataLoader(ds_train, batch_size=16, num_workers=2, pin_memory=True, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:34:05.703789Z","iopub.execute_input":"2021-11-04T20:34:05.704619Z","iopub.status.idle":"2021-11-04T20:34:05.717723Z","shell.execute_reply.started":"2021-11-04T20:34:05.704581Z","shell.execute_reply":"2021-11-04T20:34:05.716757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plot simages and mask from dataloader\nbatch = next(iter(dl_train))\nimages, masks = batch\nprint(f\"image shape: {images.shape},\\nmask shape:{masks.shape},\\nbatch len: {len(batch)}\")\n\nplt.figure(figsize=(10, 5))\n\nplt.subplot(1, 3, 1)\nplt.imshow(images[1][1])\nplt.title('Original image')\n\nplt.subplot( 1, 3, 2)\nplt.imshow(masks[1][0])\nplt.title('Mask')\n\nplt.subplot( 1, 3, 3)\nplt.imshow(images[1][1])\nplt.imshow(masks[1][0],alpha=0.2)\nplt.title('Both')\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:19.317942Z","iopub.execute_input":"2021-11-04T20:31:19.318579Z","iopub.status.idle":"2021-11-04T20:31:33.705825Z","shell.execute_reply.started":"2021-11-04T20:31:19.318535Z","shell.execute_reply":"2021-11-04T20:31:33.70508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## U-Net model architecture \n\nA u-net is commonly used for biological image segmentation because its shape allows for local and global features to be combined to create highly-precise segmentations.\n\nA u-net is shaped like an autoencoder, it has:\n\n\n\n1.   a standard convolutional network with downsampling, like one used for imagenet\n2.   upsampling layers that ultimately return an image at the same size as the input image In addition to these downsampling and upsampling blocks, it has skip connections from the downsampling blocks TO the upsampling blocks, which allows it to propagate more precise local information to the later layers.\n\n**Model is based on https://deeplearning.neuromatch.io/projects/Neuroscience/cellular_segmentation.html**","metadata":{}},{"cell_type":"code","source":"\ndef convbatchrelu(in_channels, out_channels, sz):\n  return nn.Sequential(\n      nn.Conv2d(in_channels, out_channels, sz, padding=sz//2),\n      nn.BatchNorm2d(out_channels, eps=1e-5),\n      nn.ReLU(inplace=True),\n      )\n\n\nclass convdown(nn.Module):\n  def __init__(self, in_channels, out_channels, kernel_size):\n    super().__init__()\n    self.conv = nn.Sequential()\n    for t in range(2):\n      if t == 0:\n        self.conv.add_module('conv_%d'%t,\n                             convbatchrelu(in_channels,\n                                           out_channels,\n                                           kernel_size))\n      else:\n        self.conv.add_module('conv_%d'%t,\n                             convbatchrelu(out_channels,\n                                           out_channels,\n                                           kernel_size))\n\n  def forward(self, x):\n    x = self.conv[0](x)\n    x = self.conv[1](x)\n    return x\n\n\nclass downsample(nn.Module):\n  def __init__(self, nbase, kernel_size):\n    super().__init__()\n    self.down = nn.Sequential()\n    self.maxpool = nn.MaxPool2d(2, 2)\n    for n in range(len(nbase) - 1):\n      self.down.add_module('conv_down_%d'%n,\n                           convdown(nbase[n],\n                                    nbase[n + 1],\n                                    kernel_size))\n\n  def forward(self, x):\n    xd = []\n    for n in range(len(self.down)):\n      if n > 0:\n        y = self.maxpool(xd[n - 1])\n      else:\n        y = x\n      xd.append(self.down[n](y))\n    return xd\n\n\nclass convup(nn.Module):\n  def __init__(self, in_channels, out_channels, kernel_size):\n    super().__init__()\n    self.conv = nn.Sequential()\n    self.conv.add_module('conv_0', convbatchrelu(in_channels,\n                                                 out_channels,\n                                                 kernel_size))\n    self.conv.add_module('conv_1', convbatchrelu(out_channels,\n                                                 out_channels,\n                                                 kernel_size))\n\n  def forward(self, x, y):\n    x = self.conv[0](x)\n    x = self.conv[1](x + y)\n    return x\n\n\nclass upsample(nn.Module):\n  def __init__(self, nbase, kernel_size):\n    super().__init__()\n    self.upsampling = nn.Upsample(scale_factor=2, mode='nearest')\n    self.up = nn.Sequential()\n    for n in range(len(nbase) - 1 , 0, -1):\n      self.up.add_module('conv_up_%d'%(n - 1),\n              convup(nbase[n], nbase[n - 1], kernel_size))\n\n  def forward(self, xd):\n    x = xd[-1]\n    for n in range(0, len(self.up)):\n      if n > 0:\n        x = self.upsampling(x)\n      x = self.up[n](x, xd[len(xd) - 1 - n])\n    return x\n\n\nclass Unet(nn.Module):\n  def __init__(self, nbase, nout, kernel_size):\n    super(Unet, self).__init__()\n    self.nbase = nbase\n    self.nout = nout\n    self.kernel_size = kernel_size\n    self.downsample = downsample(nbase, kernel_size)\n    nbaseup = nbase[1:]\n    nbaseup.append(nbase[-1])\n    self.upsample = upsample(nbaseup, kernel_size)\n    self.output = nn.Conv2d(nbase[1], self.nout, kernel_size,\n                            padding=kernel_size//2)\n\n  def forward(self, data):\n    T0 = self.downsample(data)\n    T0 = self.upsample(T0)\n    T0 = self.output(T0)\n    return T0\n\n  def save_model(self, filename):\n    torch.save(self.state_dict(), filename)\n\n  def load_model(self, filename, cpu=False):\n    if not cpu:\n      self.load_state_dict(torch.load(filename))\n    else:\n      self.__init__(self.nbase,\n                    self.nout,\n                    self.kernel_size,\n                    self.concatenation)\n\n      self.load_state_dict(torch.load(filename,\n                                      map_location=torch.device('cpu')))","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:33.706937Z","iopub.execute_input":"2021-11-04T20:31:33.707298Z","iopub.status.idle":"2021-11-04T20:31:33.734746Z","shell.execute_reply.started":"2021-11-04T20:31:33.70727Z","shell.execute_reply":"2021-11-04T20:31:33.73408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define the network","metadata":{}},{"cell_type":"code","source":"kernel_size = 3\nnbase = [3, 32, 64, 128, 256]  # number of channels per layer\nnout = 1  # number of outputs\n\nnet = Unet(nbase, nout, kernel_size)\n# put on GPU here if you have it\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\nnet.to(device);  # remove semi-colon to see net structure\nprint(f\"The device is {device}!!\")","metadata":{"execution":{"iopub.status.busy":"2021-11-04T21:02:32.539587Z","iopub.execute_input":"2021-11-04T21:02:32.54043Z","iopub.status.idle":"2021-11-04T21:02:32.5827Z","shell.execute_reply.started":"2021-11-04T21:02:32.540377Z","shell.execute_reply":"2021-11-04T21:02:32.581808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train the network","metadata":{}},{"cell_type":"markdown","source":"### Loss Function","metadata":{}},{"cell_type":"code","source":"def dice_loss(input, target):\n    input = torch.sigmoid(input)\n    smooth = 1.0\n    iflat = input.view(-1)\n    tflat = target.view(-1)\n    intersection = (iflat * tflat).sum()\n    return ((2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))\n\n\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma):\n        super().__init__()\n        self.gamma = gamma\n\n    def forward(self, input, target):\n        if not (target.size() == input.size()):\n            raise ValueError(\"Target size ({}) must be the same as input size ({})\"\n                             .format(target.size(), input.size()))\n        max_val = (-input).clamp(min=0)\n        loss = input - input * target + max_val + \\\n            ((-max_val).exp() + (-input - max_val).exp()).log()\n        invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))\n        loss = (invprobs * self.gamma).exp() * loss\n        return loss.mean()\n\n\nclass MixedLoss(nn.Module):\n    def __init__(self, alpha, gamma):\n        super().__init__()\n        self.alpha = alpha\n        self.focal = FocalLoss(gamma)\n\n    def forward(self, input, target):\n        loss = self.alpha*self.focal(input, target) - torch.log(dice_loss(input, target))\n        return loss.mean()","metadata":{"execution":{"iopub.status.busy":"2021-11-04T20:31:33.801157Z","iopub.execute_input":"2021-11-04T20:31:33.801881Z","iopub.status.idle":"2021-11-04T20:31:33.814001Z","shell.execute_reply.started":"2021-11-04T20:31:33.801847Z","shell.execute_reply":"2021-11-04T20:31:33.813074Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datetime import datetime\nfrom tqdm import tqdm\n\n# train the network\n# parameters related to training the network\n### you will want to increase n_epochs!\nn_epochs = 20  # number of times to cycle through all the data during training\nlearning_rate = 0.1\nweight_decay = 1e-5 # L2 regularization of weights\nmomentum = 0.9 # how much to use previous gradient direction\nn_epochs_per_save = 5 # how often to save the network\nval_frac = 0.05 # what fraction of data to use for validation\n\n# where to save the network\n# make sure to clean these out every now and then, as you will run out of space\nnow = datetime.now()\ntimestamp = now.strftime('%Y%m%dT%H%M%S')\nn_train=72\n\n\n# gradient descent flavor\noptimizer = torch.optim.SGD(net.parameters(),\n                            lr=learning_rate,\n                            weight_decay=weight_decay,\n                            momentum=0.9)\n# set learning rate schedule\nLR = np.linspace(0, learning_rate, 10)\nLR = np.append(LR, learning_rate*np.ones(n_epochs-5))\nfor i in range(5):\n    LR = np.append(LR, LR[-1]/2 * np.ones(10))\n\ncriterion = MixedLoss(10.0, 2.0)\n\n# store loss per epoch\nepoch_losses = np.zeros(n_epochs)\nepoch_losses[:] = np.nan\n\n# when we last saved the network\nsaveepoch = None\n\n# loop through entire training data set nepochs times\nfor epoch in range(n_epochs):\n  net.train() # put in train mode (affects batchnorm)\n  epoch_loss = 0\n  iters = 0\n  for param_group in optimizer.param_groups:\n    param_group['lr'] = LR[epoch]\n  with tqdm(total=545, desc=f\"Epoch {epoch + 1}/{n_epochs}\", unit='img') as pbar:\n    # loop through each batch in the training data\n    for batch_idx, batch in enumerate(dl_train):\n      # transfer to torch + GPU\n      images, masks = batch\n\n      # transfer to torch + GPU\n      images = images.to(device=device)\n      masks = masks.to(device=device)\n\n      # compute the loss\n      y = net(images)\n      loss = criterion(y, masks)\n      epoch_loss += loss.item()\n      pbar.set_postfix(**{'loss (batch)': loss.item()})\n      # gradient descent\n      optimizer.zero_grad()\n      loss.backward()\n      optimizer.step()\n      iters+=1\n      pbar.update(masks.shape[0])\n\n    epoch_losses[epoch] = epoch_loss\n    pbar.set_postfix(**{'loss (epoch)': epoch_loss})  #.update('loss (epoch) = %f'%epoch_loss)\n\n  # save checkpoint networks every now and then\n  if epoch % n_epochs_per_save == 0:\n    print(f\"\\nSaving network state at epoch {epoch+1}\")\n    saveepoch = epoch\n    savefile = f\"unet_epoch{saveepoch+1}.pth\"\n    net.save_model(savefile)\nprint(f\"\\nSaving network state at epoch {epoch+1}\")\nnet.save_model(f\"unet_epoch{epoch+1}.pth\")","metadata":{"execution":{"iopub.status.busy":"2021-11-04T21:20:12.282733Z","iopub.execute_input":"2021-11-04T21:20:12.283029Z","iopub.status.idle":"2021-11-04T21:50:10.744056Z","shell.execute_reply.started":"2021-11-04T21:20:12.282998Z","shell.execute_reply":"2021-11-04T21:50:10.742833Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_val = CellDataset(df_train, train=False)\ndl_val = DataLoader(ds_val, batch_size=4, num_workers=2, pin_memory=True, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2021-11-04T21:50:10.746778Z","iopub.execute_input":"2021-11-04T21:50:10.7473Z","iopub.status.idle":"2021-11-04T21:50:10.762506Z","shell.execute_reply.started":"2021-11-04T21:50:10.747246Z","shell.execute_reply":"2021-11-04T21:50:10.761398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def post_process(probability, threshold=0.5, min_size=20):\n    mask = cv2.threshold(probability_mask, 0.25, 1,  cv2.THRESH_BINARY)[1]\n    num_component, component = cv2.connectedComponents(mask.astype(np.uint8))\n    predictions = []\n    im = np.zeros((240, 240), np.float32)\n    for c in range(1, num_component):\n        p = (component == c)\n        if p.sum() > min_size:\n            a_prediction = np.zeros((240, 240), np.float32)\n            a_prediction[p] = 1\n            predictions.append(a_prediction)\n        im[p] = 1\n    return predictions, im","metadata":{"execution":{"iopub.status.busy":"2021-11-04T21:51:02.814664Z","iopub.execute_input":"2021-11-04T21:51:02.815013Z","iopub.status.idle":"2021-11-04T21:51:02.822797Z","shell.execute_reply.started":"2021-11-04T21:51:02.814962Z","shell.execute_reply":"2021-11-04T21:51:02.821708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"net.eval()\n\nsubmission = []\nfor i, batch in enumerate(tqdm(dl_val)):\n    preds = torch.sigmoid(net(batch[0].cuda()))\n    preds = preds.detach().cpu().numpy()[:, 0, :, :] # (batch_size, 1, size, size) -> (batch_size, size, size)\n    for index,probability_mask in  enumerate(preds):\n            print(f\"\\nsum prob: {np.sum(probability_mask)}\")\n            plt.figure(figsize=(10, 5))\n            probability_mask = cv2.resize(probability_mask, dsize=(240, 240), interpolation=cv2.INTER_LINEAR)\n            plt.subplot(1, 3, 1)\n            plt.imshow(batch[1][index][0])\n            plt.title('Original Mask')\n\n            plt.subplot( 1, 3, 2)\n            plt.imshow(probability_mask)\n            plt.title('Probality Prediction')\n\n            predictions, pred_mask = post_process(probability_mask, threshold=0.65)\n            print(f\"prediction num {len(predictions)}\")\n            plt.subplot(1, 3, 3)\n            plt.imshow(pred_mask)\n            plt.title('Predicted Mask')\n            plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-04T21:51:15.395421Z","iopub.execute_input":"2021-11-04T21:51:15.396284Z","iopub.status.idle":"2021-11-04T21:51:54.372777Z","shell.execute_reply.started":"2021-11-04T21:51:15.396234Z","shell.execute_reply":"2021-11-04T21:51:54.371801Z"},"trusted":true},"execution_count":null,"outputs":[]}]}