{"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 with Deep watershed transform(DWT) [Infer]\n[[Train notebook]](https://www.kaggle.com/ebinan92/unet-with-deep-watershed-transform-dwt-train)  \nInference pipeline is almost same as [Awsaf's notebook](https://www.kaggle.com/awsaf49/pytorch-sartorius-unet-strikes-back-infer) expect watershed algorithm added. ","metadata":{}},{"cell_type":"markdown","source":"### import, seed, config","metadata":{}},{"cell_type":"code","source":"!pip install -q ../input/pytorch-segmentation-models-lib/pretrainedmodels-0.7.4/pretrainedmodels-0.7.4\n!pip install -q ../input/pytorch-segmentation-models-lib/efficientnet_pytorch-0.6.3/efficientnet_pytorch-0.6.3\n!pip install -q ../input/pytorch-segmentation-models-lib/timm-0.4.12-py3-none-any.whl\n!pip install -q ../input/pytorch-segmentation-models-lib/segmentation_models_pytorch-0.2.0-py3-none-any.whl","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T08:17:34.978936Z","iopub.execute_input":"2021-11-26T08:17:34.979679Z","iopub.status.idle":"2021-11-26T08:18:08.193249Z","shell.execute_reply.started":"2021-11-26T08:17:34.979642Z","shell.execute_reply":"2021-11-26T08:18:08.192363Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import skimage.morphology\nimport segmentation_models_pytorch as smp\nimport cupy as cp\nimport os\nimport skimage\nfrom skimage.morphology import thin\nfrom scipy import ndimage as ndi\nfrom skimage.measure import label\nfrom skimage.segmentation import watershed\nimport numpy as np\nfrom albumentations.pytorch import ToTensorV2\nimport albumentations as A\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nimport torch\nimport matplotlib.pyplot as plt\nimport cv2\nimport gc\nfrom tqdm import tqdm\nfrom glob import glob\nimport random\nimport pandas as pd","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T08:18:08.197149Z","iopub.execute_input":"2021-11-26T08:18:08.19738Z","iopub.status.idle":"2021-11-26T08:18:08.206947Z","shell.execute_reply.started":"2021-11-26T08:18:08.197353Z","shell.execute_reply":"2021-11-26T08:18:08.206283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class config:\n    SAMPLE_SUBMISSION = '../input/sartorius-cell-instance-segmentation/sample_submission.csv'\n    TRAIN_CSV = \"../input/sartorius-cell-instance-segmentation/train.csv\"\n    TRAIN_PATH = \"../input/sartorius-cell-instance-segmentation/train\"\n    TEST_PATH = \"../input/sartorius-cell-instance-segmentation/test\"\n    MODEL_PATH = \"../input/resnet101-dwt/models\"\n    RESNET_MEAN = (0.485, 0.456, 0.406)\n    RESNET_STD = (0.229, 0.224, 0.225)\n    IMAGE_RESIZE = (512, 704)\n    model_name = 'resnet101'\n    device = 'cuda'\n    BS = 1\n    num_workers = 2\n    ttas = [0, 1, 2, 3]\n    mask_len = 6\n    \n\ndef fix_all_seeds(seed):\n    np.random.seed(seed)\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\n\nfix_all_seeds(2021)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T08:18:37.13927Z","iopub.execute_input":"2021-11-26T08:18:37.139539Z","iopub.status.idle":"2021-11-26T08:18:37.150836Z","shell.execute_reply.started":"2021-11-26T08:18:37.139509Z","shell.execute_reply":"2021-11-26T08:18:37.150109Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.DataFrame(glob(f'{config.TEST_PATH}/*'), columns=['image_path'])\ntest_df['id'] = test_df.image_path.map(lambda x: x.split('/')[-1].split('.')[0])","metadata":{"execution":{"iopub.status.busy":"2021-11-26T08:18:37.635964Z","iopub.execute_input":"2021-11-26T08:18:37.636536Z","iopub.status.idle":"2021-11-26T08:18:37.659664Z","shell.execute_reply.started":"2021-11-26T08:18:37.636499Z","shell.execute_reply":"2021-11-26T08:18:37.658911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dataset and Augmentation","metadata":{}},{"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, df, transforms=None):\n        self.df = df\n        self.img_paths = df['image_path'].values\n        try:  # if there is no mask then only send images --> test data\n            self.msk_paths = df['mask_path'].values\n        except BaseException:\n            self.msk_paths = None\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, index):\n        img_path = self.img_paths[index]\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        if self.msk_paths is not None:\n            msk_path = self.msk_paths[index]\n            msk = np.load(msk_path)\n            if self.transforms:\n                data = self.transforms(image=img, mask=msk)\n                img = data['image']\n                msk = data['mask']\n            msk = np.expand_dims(msk, axis=0)  # output_shape: (batch_size, 1, img_size, img_size)\n            return img, msk\n        else:\n            if self.transforms:\n                data = self.transforms(image=img)\n                img = data['image']\n            return img, img_path\n        \ndata_transforms = {\n    \"valid\": A.Compose([\n        A.Resize(config.IMAGE_RESIZE[0], config.IMAGE_RESIZE[1]),\n        A.Normalize(mean=config.RESNET_MEAN, std=config.RESNET_STD, p=1),\n        ToTensorV2()], p=1.0)\n}","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T08:18:37.938224Z","iopub.execute_input":"2021-11-26T08:18:37.938797Z","iopub.status.idle":"2021-11-26T08:18:37.949705Z","shell.execute_reply.started":"2021-11-26T08:18:37.938756Z","shell.execute_reply":"2021-11-26T08:18:37.948527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Utils","metadata":{}},{"cell_type":"code","source":"def ins2rle(ins):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    ins = cp.array(ins)\n    pixels = ins.flatten()\n    pad = cp.array([0])\n    pixels = cp.concatenate([pad, pixels, pad])\n    runs = cp.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n\ndef mask2rle(lab_mask, cutoff=0.5, min_object_size=1.0):\n    \"\"\" Return run length encoding of mask.\n        ref: https://www.kaggle.com/raoulma/nuclei-dsb-2018-tensorflow-u-net-score-0-352\n    \"\"\"\n    # segment image and label different objects\n#     lab_mask = skimage.morphology.label(mask > cutoff)\n\n    # Keep only objects that are large enough.\n    (mask_labels, mask_sizes) = np.unique(lab_mask, return_counts=True)\n    if (mask_sizes < min_object_size).any():\n        mask_labels = mask_labels[mask_sizes < min_object_size]\n        for n in mask_labels:\n            lab_mask[lab_mask == n] = 0\n        lab_mask = skimage.morphology.label(lab_mask > cutoff)\n\n    # Loop over each object excluding the background labeled by 0.\n    for i in range(1, lab_mask.max() + 1):\n        yield ins2rle(lab_mask == i)\n\n\ndef aug(img, axis=0):\n    if axis == 1:\n        return torch.flip(img, dims=(1,))\n    elif axis == 2:\n        return torch.flip(img, dims=(2,))\n    elif axis == 3:\n        return torch.flip(img, dims=(1, 2))\n    elif axis == 4:\n        return torch.rot90(img, k=1, dims=(1, 2))\n    elif axis == 5:\n        return torch.rot90(img, k=1, dims=(2, 1))\n    else:\n        return img\n\n\ndef reverse_aug(img, axis=0):\n    if axis == 1:\n        return torch.flip(img, dims=(1,))\n    elif axis == 2:\n        return torch.flip(img, dims=(2,))\n    elif axis == 3:\n        return torch.flip(img, dims=(1, 2))\n    elif axis == 4:\n        return torch.rot90(img, k=1, dims=(2, 1))\n    elif axis == 5:\n        return torch.rot90(img, k=1, dims=(1, 2))\n    else:\n        return img\n\n\ndef get_aug_img(img, ttas=config.ttas):\n    \"\"\"\n    Args:\n        img  :  image\n        ttas :  tta modes ex [0, 1]\n    Return:\n        augmentated images shape (num_tta, dim0, dim1, channel)\n    \"\"\"\n    if len(ttas) == 0:\n        return img.unsqueeze(0)\n    aug_img = []\n    for idx, tta_mode in enumerate(ttas):\n        aug_img.append(aug(img, axis=tta_mode))\n    aug_img = torch.stack(aug_img, dim=0)\n    return aug_img\n\n\ndef fix_aug_img(aug_pred, ttas=config.ttas):\n    \"\"\"\n    Args:\n        aug_pred  :  prediction of augmented images\n        ttas      :  tta modes ex [0, 1]\n    Return:\n        final image after ensemble\n    \"\"\"\n    if len(ttas) == 0:\n        return aug_pred\n    fixed_pred = []\n    for idx, tta_mode in enumerate(ttas):\n        fixed_pred.append(reverse_aug(aug_pred[idx], axis=tta_mode))\n    fixed_pred = torch.stack(fixed_pred, dim=0)\n    fixed_pred = torch.mean(fixed_pred, dim=0)\n    return fixed_pred\n\n\ndef watershed_energy(msk=None,\n                     energy=None,\n                     threshold=0.5,\n                     threshold_energy=0.6,\n                     line=False):\n\n    msk_ths = (np.copy(msk) > 255 * threshold) * 1\n    energy_ths = (np.copy(energy) > 255 * threshold_energy) * 1\n\n    markers = label(energy_ths)\n    labels = watershed(-energy,\n                       markers,\n                       mask=msk_ths,\n                       watershed_line=line)\n    return labels\n\n\ndef load_model(path):\n    model = smp.Unet(config.model_name, encoder_weights=None, activation=None, classes=config.mask_len)\n    model.load_state_dict(torch.load(path))\n    model = model.to(config.device)\n    model.eval()\n    return model","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T08:18:38.629525Z","iopub.execute_input":"2021-11-26T08:18:38.62999Z","iopub.status.idle":"2021-11-26T08:18:38.652287Z","shell.execute_reply.started":"2021-11-26T08:18:38.629949Z","shell.execute_reply":"2021-11-26T08:18:38.651493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Inference","metadata":{}},{"cell_type":"code","source":"@torch.no_grad()\ndef infer(model_paths, test_loader, num_log=3):\n    pred_strings = []\n    pred_paths = []\n    msks = []\n    imgs = []\n    energys = []\n    for idx, (img, img_path) in enumerate(tqdm(test_loader, total=len(test_loader), desc='Infer')):\n        img = img.to(config.device, dtype=torch.float).squeeze()\n        img = get_aug_img(img, ttas=config.ttas)\n        msk = []\n        energy = []\n        for path in model_paths:\n            model = load_model(path)\n            out = model(img).squeeze(0)  # removing batch axis\n            out = fix_aug_img(out, ttas=config.ttas)\n            out = nn.Sigmoid()(out)  # removing channel axis\n            msk.append(out[0])\n            energy.append(torch.mean(out[:-1], dim=0))\n        msk = torch.mean(torch.stack(msk, dim=0), dim=0)\n        msk = F.interpolate(msk[None, None, ], size=(520, 704), mode='nearest')[0, 0]\n        msk = msk.cpu().detach().numpy()\n        energy = torch.mean(torch.stack(energy, dim=0), dim=0)\n        energy = F.interpolate(energy[None, None, ], size=(520, 704), mode='nearest')[0, 0]\n        energy = energy.cpu().detach().numpy()\n        img = F.interpolate(img[0:1, ], size=(520, 704), mode='nearest')[0]  # first dim is image w/o aug\n        img = img.squeeze().permute((1, 2, 0)).cpu().detach().numpy()\n        msk = watershed_energy(msk * 255, energy * 255, 0.5, 0.7)\n        if idx < num_log:\n            msks.append(msk)\n            energys.append(energy)\n            imgs.append(img)\n        rle = list(mask2rle(msk))\n        pred_strings.extend(rle)\n        pred_paths.extend(img_path * len(rle))\n        del img, msk\n        gc.collect()\n        torch.cuda.empty_cache()\n    return pred_strings, pred_paths, imgs, msks, energys","metadata":{"execution":{"iopub.status.busy":"2021-11-26T08:18:39.248055Z","iopub.execute_input":"2021-11-26T08:18:39.248632Z","iopub.status.idle":"2021-11-26T08:18:39.262324Z","shell.execute_reply.started":"2021-11-26T08:18:39.248594Z","shell.execute_reply":"2021-11-26T08:18:39.261396Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = TestDataset(test_df, transforms=data_transforms['valid'])\ntest_loader = DataLoader(test_dataset, batch_size=config.BS,\n                         num_workers=config.num_workers, shuffle=False, pin_memory=True)\nmodel_paths = glob(f'{config.MODEL_PATH}/resnet101_*.pth')\n\npred_strings, pred_paths, imgs, msks, energys = infer(model_paths, test_loader)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T08:18:39.47175Z","iopub.execute_input":"2021-11-26T08:18:39.472208Z","iopub.status.idle":"2021-11-26T08:19:23.462453Z","shell.execute_reply.started":"2021-11-26T08:18:39.472169Z","shell.execute_reply":"2021-11-26T08:19:23.461658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Check result","metadata":{}},{"cell_type":"code","source":"for img, msk, energy in zip(imgs, msks, energys):\n    plt.figure(figsize=(15, 7))\n    plt.subplot(1, 3, 1)\n    plt.imshow(img)\n    plt.axis('OFF')\n    plt.title('image')\n    plt.subplot(1, 3, 2)\n    plt.imshow(msk)\n    plt.axis('OFF')\n    plt.title('mask')\n    plt.subplot(1, 3, 3)\n    plt.imshow(energy)\n    plt.axis('OFF')\n    plt.title('energy')\n    plt.tight_layout()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T08:19:23.464565Z","iopub.execute_input":"2021-11-26T08:19:23.464818Z","iopub.status.idle":"2021-11-26T08:19:24.988623Z","shell.execute_reply.started":"2021-11-26T08:19:23.464783Z","shell.execute_reply":"2021-11-26T08:19:24.987924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Submission","metadata":{}},{"cell_type":"code","source":"ids = list(map(lambda x: x.split('/')[-1].split('.')[0], pred_paths))\npred_df = pd.DataFrame({'id': ids,\n                        'predicted': pred_strings})\nsub_df = pd.read_csv('/kaggle/input/sartorius-cell-instance-segmentation/sample_submission.csv')\ndel sub_df['predicted']\nsub_df = sub_df.merge(pred_df, on='id', how='left')\nsub_df.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T08:19:24.989892Z","iopub.execute_input":"2021-11-26T08:19:24.990447Z","iopub.status.idle":"2021-11-26T08:19:25.027411Z","shell.execute_reply.started":"2021-11-26T08:19:24.990408Z","shell.execute_reply":"2021-11-26T08:19:25.026701Z"},"trusted":true},"execution_count":null,"outputs":[]}]}