{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# UNET-Efficientnet-B0\n\nThis notebook uses a UNET architecture with an EfficientNet-B0 backbone with pretrained weights.\n\n**Reference Notebook:** https://www.kaggle.com/julian3833/sartorius-starter-baseline-torch-u-net","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport time\nimport random\nimport numpy as np\nimport pandas as pd\nimport plotly.express as px\nfrom matplotlib import pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts, CosineAnnealingLR\n\n# !pip install -U git+https://github.com/albumentations-team/albumentations\n!pip install ../input/segmentation-libs/albumentations-1.1.0/albumentations-1.1.0\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n!pip install ../input/segmentation-libs/timm-0.3.2-py3-none-any.whl\n!pip install ../input/segmentation-libs/pretrainedmodels-0.7.4-py3-none-any.whl\n!pip install ../input/segmentation-libs/efficientnet_pytorch-0.6.3-py3-none-any.whl\n!pip install ../input/segmentation-libs/segmentation_models_pytorch-0.1.3-py3-none-any.whl\n\nimport collections.abc as container_abcs\ntorch._six.container_abcs = container_abcs\nimport segmentation_models_pytorch as smp","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# EDA","metadata":{}},{"cell_type":"code","source":"train = \"../input/sartorius-cell-instance-segmentation/train\"\ntest = \"../input/sartorius-cell-instance-segmentation/test\"\nsave_path = os.getcwd()\n\ntrain_df = pd.read_csv(os.path.join(os.path.dirname(train), \"train.csv\"))\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_ids = train_df[\"id\"].unique()\nprint(f\"There are {len(image_ids)} images in the dataset\")\ncell_instances = len(train_df[\"annotation\"])\nprint(f\"There are {cell_instances} cell instances in the dataset\")\ncell_types = train_df[\"cell_type\"].unique()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"px.histogram(train_df[\"cell_type\"], labels=cell_types)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMAGE_SHAPE = (train_df[\"height\"][0], train_df[\"width\"][0])\nIMAGE_RESIZE = (224, 224)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_decode(mask_rle):\n    \n    mask_rle = np.array(mask_rle.split(), dtype=np.int)\n    pixels = mask_rle.reshape(-1, 2)\n#     assert len(start) == len(length)\n    pixels[:, 0] -= 1\n    mask = np.zeros(IMAGE_SHAPE[0] * IMAGE_SHAPE[1])\n    for pixel in pixels:\n        mask[pixel[0]:pixel[0] + pixel[1]] = 1\n    return mask.reshape(IMAGE_SHAPE)\n\ndef prepare_image_mask(mask_annotations):\n    \n    mask = np.zeros(IMAGE_SHAPE)\n    for mask_annotation in mask_annotations:\n        mask += rle_decode(mask_annotation)\n        \n    mask = mask.clip(0, 1)\n    return mask\n\ndef compute_iou(labels, y_pred):\n    \"\"\"\n    Computes the IoU for instance labels and predictions.\n\n    Args:\n        labels (np array): Labels.\n        y_pred (np array): predictions\n\n    Returns:\n        np array: IoU matrix, of size true_objects x pred_objects.\n    \"\"\"\n\n    true_objects = len(np.unique(labels))\n    pred_objects = len(np.unique(y_pred))\n\n    # Compute intersection between all objects\n    intersection = np.histogram2d(\n        labels.flatten(), y_pred.flatten(), bins=(true_objects, pred_objects)\n    )[0]\n\n    # Compute areas (needed for finding the union between all objects)\n    area_true = np.histogram(labels, bins=true_objects)[0]\n    area_pred = np.histogram(y_pred, bins=pred_objects)[0]\n    area_true = np.expand_dims(area_true, -1)\n    area_pred = np.expand_dims(area_pred, 0)\n\n    # Compute union\n    union = area_true + area_pred - intersection\n    iou = intersection / union\n    \n    return iou[1:, 1:]  # exclude background\n\ndef precision_at(threshold, iou):\n    \"\"\"\n    Computes the precision at a given threshold.\n\n    Args:\n        threshold (float): Threshold.\n        iou (np array): IoU matrix.\n\n    Returns:\n        int: Number of true positives,\n        int: Number of false positives,\n        int: Number of false negatives.\n    \"\"\"\n    matches = iou > threshold\n    true_positives = np.sum(matches, axis=1) == 1  # Correct objects\n    false_positives = np.sum(matches, axis=0) == 0  # Missed objects\n    false_negatives = np.sum(matches, axis=1) == 0  # Extra objects\n    tp, fp, fn = (\n        np.sum(true_positives),\n        np.sum(false_positives),\n        np.sum(false_negatives),\n    )\n    return tp, fp, fn\n\ndef iou_map(truths, preds, verbose=0):\n    \"\"\"\n    Computes the metric for the competition.\n    Masks contain the segmented pixels where each object has one value associated,\n    and 0 is the background.\n\n    Args:\n        truths (list of masks): Ground truths.\n        preds (list of masks): Predictions.\n        verbose (int, optional): Whether to print infos. Defaults to 0.\n\n    Returns:\n        float: mAP.\n    \"\"\"\n    ious = [compute_iou(truth, pred) for truth, pred in zip(truths, preds)]\n\n    if verbose:\n        print(\"Thresh\\tTP\\tFP\\tFN\\tPrec.\")\n\n    prec = []\n    for t in np.arange(0.5, 1.0, 0.05):\n        tps, fps, fns = 0, 0, 0\n        for iou in ious:\n            tp, fp, fn = precision_at(t, iou)\n            tps += tp\n            fps += fp\n            fns += fn\n\n        p = tps / (tps + fps + fns)\n        prec.append(p)\n\n        if verbose:\n            print(\"{:1.3f}\\t{}\\t{}\\t{}\\t{:1.3f}\".format(t, tps, fps, fns, p))\n\n    if verbose:\n        print(\"AP\\t-\\t-\\t-\\t{:1.3f}\".format(np.mean(prec)))\n\n    return np.mean(prec)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"paths = []\nfor image_id in image_ids:\n    paths.append(os.path.join(train, image_id + \".png\"))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"idx = np.random.choice(np.arange(len(paths)))\nimg_path = paths[idx]\nimg = cv2.imread(img_path)\nmask_annotations = train_df[\"annotation\"][train_df[\"id\"] == image_ids[idx]].values\nprint(f\"Chose image {img_path}.\")\nprint(f\"This image has {len(mask_annotations[0:][::2])} cells.\")\n\nmask = prepare_image_mask(mask_annotations)\nmask = np.clip(mask, 0, 1)\nfig, ax = plt.subplots(1, 3, figsize=(20, 15))\nfig.suptitle(f\"Cell Type: {train_df['cell_type'][train_df['id'] == image_ids[idx]].values[0]}\")\nax[0].imshow(img, cmap=\"bone_r\")\nax[1].imshow(img, cmap=\"bone_r\")\nax[1].imshow(mask, alpha=0.2, cmap=\"bone_r\")\nax[2].imshow(mask, cmap=\"gray\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class SartoriusDataset(Dataset):\n    \n    def __init__(self, image_paths, data_df, transforms=None):\n        self.image_paths = image_paths\n        self.data_df = data_df\n        self.transforms = transforms\n    \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        image_path = self.image_paths[idx]\n        image_id = image_path.split('/')[-1].split('.')[0]\n        image = cv2.imread(image_path)\n        \n        image_annotations = self.data_df[\"annotation\"][self.data_df[\"id\"] == image_id].values\n        mask = prepare_image_mask(image_annotations)\n        mask = (mask >= 1).astype(np.float32)\n        if self.transforms is not None:\n            augmented = self.transforms(image=image, mask=mask)\n            image = augmented[\"image\"]\n            mask = augmented[\"mask\"]\n\n        return image, mask.unsqueeze(0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transforms_train = A.Compose([\n  A.RandomResizedCrop(IMAGE_RESIZE[0], IMAGE_RESIZE[1], scale=(0.9, 1), p=1), \n  A.HorizontalFlip(p=0.5),\n  A.ShiftScaleRotate(p=0.5),\n  A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, p=0.7),\n  A.RandomBrightnessContrast(brightness_limit=(-0.2,0.2), contrast_limit=(-0.2, 0.2), p=0.7),\n  A.CLAHE(clip_limit=(1,4), p=0.5),\n  A.OneOf([\n    A.OpticalDistortion(distort_limit=1.0),\n    A.GridDistortion(num_steps=5, distort_limit=1.),\n    A.ElasticTransform(alpha=3),\n  ], p=0.2),\n  A.OneOf([\n    A.GaussNoise(var_limit=[10, 50]),\n    A.GaussianBlur(),\n    A.MotionBlur(),\n    A.MedianBlur(),\n  ], p=0.2),\n  A.Resize(IMAGE_RESIZE[0], IMAGE_RESIZE[1]),\n  A.OneOf([\n    A.JpegCompression(),\n    A.Downscale(scale_min=0.1, scale_max=0.15),\n    ], p=0.2),\n  A.IAAPiecewiseAffine(p=0.2),\n  A.IAASharpen(p=0.2),\n  A.Cutout(max_h_size=int(IMAGE_RESIZE[0] * 0.1), max_w_size=int(IMAGE_RESIZE[1] * 0.1), num_holes=5, p=0.5),\n  A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n  A.VerticalFlip(p=0.5),\n  ToTensorV2()\n])\n\ntransforms_valid = A.Compose([\n  A.Resize(IMAGE_RESIZE[0], IMAGE_RESIZE[1]),\n  A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n  ToTensorV2()\n])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = SartoriusDataset(paths, train_df, transforms = transforms_train)\n\nfor i in range(5):\n    image, mask = train_dataset[i]\n    print(f\"Image size: {image.size()}, Mask size: {mask.size()}\")\n    plt.imshow(image.permute(1, 2, 0))\n    plt.show()\n    plt.imshow(mask[0], cmap=\"gray\")\n    plt.show() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def dice_loss(pred, target):\n    pred = pred.sigmoid().view(-1)\n    target = target.view(-1)\n    \n    numerator = 2.0 * (pred * target).sum() + 1.0\n    denominator = pred.sum() + target.sum() + 1.0\n    \n    return numerator / denominator\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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomUnet(nn.Module):\n    \n    def __init__(self, model_name=\"efficientnet-b0\"):\n        super(CustomUnet, self).__init__()\n        self.model = smp.Unet(model_name, encoder_weights=\"imagenet\", in_channels=3, classes=1, activation=None)\n        \n    def forward(self, x):\n        prediction = self.model(x)\n        return prediction","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(model, train_loader, criterion, optimizer, device):\n    model.train()\n    losses = []\n\n    for batch_idx, (images, masks) in enumerate(train_loader):\n\n        images, masks = images.to(device), masks.to(device)\n        optimizer.zero_grad()\n        logits = model(images)\n        loss = criterion(logits, masks)\n        loss.backward()\n        optimizer.step()\n\n        losses.append(loss.item())\n        if batch_idx % 5 == 0:\n            print(f'Training Loss at {batch_idx}th batch: {loss.item():.4f}')\n\n    loss_train = np.mean(losses)\n    return loss_train\n\ndef val_fn(model, val_loader, criterion, device):\n    model.eval()\n    losses = []\n    trues = []\n    preds = []\n\n    with torch.no_grad():\n        for batch_idx, (images, masks) in enumerate(val_loader):\n\n            images, masks = images.to(device), masks.to(device)\n\n            logits = model(images)\n            loss = criterion(logits, masks)\n            \n            preds.append(logits.sigmoid().detach().cpu().numpy())\n            trues.append(masks.detach().cpu().numpy())\n            losses.append(loss.item())\n            if batch_idx % 5 == 0:\n                print(f'Validation Loss at {batch_idx}th batch: {loss.item():.4f}')\n\n    loss_valid = np.mean(losses)\n    return loss_valid, trues, preds\n\ndef post_process_preds(preds, threshold=0.5):\n    \"\"\"\n        Only thresholds images and returns a list of images.\n    \"\"\"\n    processed_preds = []\n    for pred in preds:\n        pred = pred[:, 0, :, :]\n        pred = (pred > threshold).astype(np.float32)\n        processed_preds.append(pred)\n    \n    return processed_preds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_loop(epochs, lr, early_stopping, device, print_freq=1):\n    \n    # This is a very straightforward model with no folding\n    # Just do a 90/10 split for train/val.\n    idx = int(len(paths) * 0.9)\n    train_dataset = SartoriusDataset(paths[:idx], train_df, transforms=transforms_train)\n    val_dataset = SartoriusDataset(paths[idx:], train_df, transforms=transforms_valid)\n    \n    train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)\n    val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0)\n    \n    model = CustomUnet()\n    model.to(device)\n    optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-6, amsgrad=False)\n    scheduler = CosineAnnealingLR(optimizer, T_max=5)\n    \n    criterion = MixedLoss(10, 2)\n    best_loss = np.inf\n    best_avg_prec = -np.inf\n    \n    early_stop_epochs = 0\n    for epoch in range(epochs):\n        \n        start_time = time.time()\n        train_loss = train_fn(model, train_dataloader, criterion, optimizer, device)\n        \n        val_loss, trues, preds = val_fn(model, val_dataloader, criterion, device)\n        scheduler.step()\n\n        time_taken = time.time() - start_time\n        \n        preds = post_process_preds(preds)\n        avg_prec = iou_map(trues, preds, 1)\n        if epoch % print_freq == 0:\n            print(f'At Epoch {epoch}, Training Loss: {train_loss:.5f}, Validation Loss: {val_loss:.4f}, Avg Prec: {avg_prec:.4f}. Took {time_taken:.0f}s.')\n        \n        if avg_prec > best_avg_prec:\n            best_avg_prec = avg_prec\n            torch.save(model.state_dict(), os.path.join(save_path, f\"unet_{epoch}_{best_avg_prec:.4f}.pth\"))\n        \n        if val_loss < best_loss:\n            best_loss = val_loss\n            print(f\"Validation loss < Best loss. Saving model...\")\n            torch.save(model.state_dict(), os.path.join(save_path, f\"unet_{epoch}_{val_loss:.2f}.pth\"))\n            early_stop_epochs = 0\n        else:\n            early_stop_epochs += 1\n            if early_stop_epochs == early_stopping:\n                break\n    \n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir -p /root/.cache/torch/hub/checkpoints/\n!cp ../input/efficientnetv2weights/efficientnet-b0-355c32eb.pth  /root/.cache/torch/hub/checkpoints/efficientnet-b0-355c32eb.pth\n\nLR = 5e-4\nEPOCHS = 100\nEARLY_STOPPING = 15\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\nmodel = train_loop(EPOCHS, LR, EARLY_STOPPING, device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prediction","metadata":{}},{"cell_type":"markdown","source":"## Utilities","metadata":{}},{"cell_type":"code","source":"def separate_mask_components(mask, min_size=300):\n    num_component, component = cv2.connectedComponents(mask.astype(np.uint8))\n    predictions = []\n    for c in range(1, num_component):\n        p = (component == c)\n        if p.sum() > min_size:\n            a_prediction = np.zeros(IMAGE_SHAPE, np.float32)\n            a_prediction[p] = 1\n            predictions.append(a_prediction)\n    return predictions\n\ndef rle_encoding(x):\n    dots = np.where(x.flatten() == 1)[0]\n    run_lengths = []\n    prev = -2\n    for b in dots:\n        if (b>prev+1): run_lengths.extend((b + 1, 0))\n        run_lengths[-1] += 1\n        prev = b\n    return ' '.join(map(str, run_lengths))\n\ndef one_hot(y, num_classes, dtype=np.uint8):\n    y = np.array(y, dtype='int')\n    input_shape = y.shape\n    if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:\n        input_shape = tuple(input_shape[:-1])\n    y = y.ravel()\n    if not num_classes:\n        num_classes = np.max(y) + 1\n    n = y.shape[0]\n    categorical = np.zeros((n, num_classes), dtype=dtype)\n    categorical[np.arange(n), y] = 1\n    output_shape = input_shape + (num_classes,)\n    categorical = np.reshape(categorical, output_shape)\n    return categorical\n\ndef fix_overlap(msk):\n    \"\"\"\n    Args:\n        mask: multi-channel mask, each channel is an instance of cell, shape:(520,704,None)\n    Returns:\n        multi-channel mask with non-overlapping values, shape:(520,704,None)\n    \"\"\"\n    msk = np.array(msk)\n#     print(msk.shape)\n#     msk = np.pad(msk, [[0,0],[0,0],[1,0]]) # add dummy mask for background\n#     ins_len = msk.shape[-1]\n    msk = np.argmax(msk,axis=-1)# convert multi channel mask to single channel mask, argmax will remove overlap\n    msk = one_hot(msk, num_classes=ins_len) # back to multi-channel mask, some instance might get removed\n    msk = msk[...,1:] # remove background mask\n    msk = msk[...,np.any(msk, axis=(0,1))] # remove all-zero masks\n    #assert np.prod(msk, axis=-1).sum()==0 # overlap check, will raise error if there is overlap\n    return msk\n\ndef check_overlap(msk):\n    msk = msk.astype(np.bool).astype(np.uint8) # binary mask\n    return np.any(np.sum(msk, axis=-1)>1) # only one channgel will contain value","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SartoriusTestDataset(Dataset):\n    \n    def __init__(self, root, transforms=None):\n        self.root = root\n        self.image_ids = [f_name[:-4] for f_name in os.listdir(self.root)]\n        self.transforms = transforms\n        \n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        image_path = os.path.join(self.root, image_id + \".png\")\n        image = cv2.imread(image_path)\n        \n        if self.transforms is not None:\n            augmented = self.transforms(image=image)\n            image = augmented[\"image\"]\n        return image, image_id\n    \n    def __len__(self):\n        return len(self.image_ids)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = SartoriusTestDataset(test, transforms_valid)\ntest_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n\nweights = [w_file for w_file in os.listdir(save_path) if \"0.\" in w_file[:-5]]\navgs = [float(w[:-4].split('_')[-1]) for w in weights]\nbest_avg_weights = weights[np.argmax(avgs)]\nprint(best_avg_weights)\nmodel.load_state_dict(torch.load(os.path.join(save_path, best_avg_weights)))\nmodel.eval()\n\nsubmission = []\nfor image, image_id in test_dataloader:\n    prediction = model(image.to(device)).sigmoid().detach().cpu().numpy()\n    prediction = post_process_preds([prediction])[0][0]\n    probability_mask = cv2.resize(prediction, dsize=(IMAGE_SHAPE[1], IMAGE_SHAPE[0]), interpolation=cv2.INTER_LINEAR)\n    \n    cell_instances = separate_mask_components(probability_mask)\n    cell_instances = np.stack(cell_instances, axis=-1)\n    if check_overlap(cell_instances):\n        cell_instances = fix_overlap(cell_instances)\n        \n    for cell in cell_instances:\n        submission.append([image_id[0], rle_encoding(cell)])\n    \n    image_ids = [image_id for image_id, cell in submission]\n    if image_id not in image_ids:\n        submission.append([image_id, \"\"])\n\ndf = pd.DataFrame(submission, columns=[\"id\", \"predicted\"])\ndf.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}