{"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) [Train]\n[[Inference notebook]](https://www.kaggle.com/ebinan92/unet-with-deep-watershed-transform-dwt-infer)  \nThis notebook is the simple implementation of DWT method [(paper)](https://arxiv.org/abs/1611.08303). <br>\nOrignal paper's approch use two Unet to learn disntace transformation.  \nFor simplicity, I tried single Unet with multi task learning approach invented by snakers41.  \nPlease see the detail info at [his blog](https://spark-in.me/post/playing-with-dwt-and-ds-bowl-2018)","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-26T09:27:09.037803Z","iopub.execute_input":"2021-11-26T09:27:09.038597Z","iopub.status.idle":"2021-11-26T09:27:43.462154Z","shell.execute_reply.started":"2021-11-26T09:27:09.038492Z","shell.execute_reply":"2021-11-26T09:27:43.461118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm import tqdm_notebook as tqdm\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom sklearn.model_selection import StratifiedKFold\nfrom skimage.segmentation import watershed\nfrom skimage.measure import label\nimport matplotlib.pyplot as plt\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport torch\nimport segmentation_models_pytorch as smp\nimport pandas as pd\nimport numpy as np\nimport cv2\nfrom albumentations.pytorch import ToTensorV2\nimport albumentations as A\nimport random\nimport pickle\nimport os\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)\n\n\nclass 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 = \"models\"\n    MASK_PATH = \"../input/dwtmasks\"\n    RESNET_MEAN = (0.485, 0.456, 0.406)\n    RESNET_STD = (0.229, 0.224, 0.225)\n    IMAGE_RESIZE = (512, 704)\n    LR = 5e-4\n    min_LR = 5e-5\n    model_name = 'resnet101'\n    device = 'cuda'\n    BS = 8\n    num_workers = 2\n    N_EPOCH = 50\n    N_FOLD = 5\n    mask_len = 6\n    \nos.makedirs(config.MODEL_PATH, exist_ok=True)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T09:27:43.464543Z","iopub.execute_input":"2021-11-26T09:27:43.465858Z","iopub.status.idle":"2021-11-26T09:27:53.363791Z","shell.execute_reply.started":"2021-11-26T09:27:43.465811Z","shell.execute_reply":"2021-11-26T09:27:53.362836Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train = pd.read_csv(config.TRAIN_CSV).groupby('id').first().reset_index()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:27:53.365267Z","iopub.execute_input":"2021-11-26T09:27:53.365548Z","iopub.status.idle":"2021-11-26T09:27:53.941305Z","shell.execute_reply.started":"2021-11-26T09:27:53.365512Z","shell.execute_reply":"2021-11-26T09:27:53.940598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## check mask dataset","metadata":{}},{"cell_type":"code","source":"image_id = df_train.iloc[0].id\nwith open(f'{config.MASK_PATH}/{image_id}.pickle', 'rb') as f:\n    masks = pickle.load(f)\n\nlen(masks)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:27:53.94303Z","iopub.execute_input":"2021-11-26T09:27:53.94328Z","iopub.status.idle":"2021-11-26T09:27:54.176655Z","shell.execute_reply.started":"2021-11-26T09:27:53.943244Z","shell.execute_reply":"2021-11-26T09:27:54.175946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"masks[0]: default mask, masks[1:4]: thin_mask, masks[5]: border_mask ,masks[6]: instance mask","metadata":{}},{"cell_type":"code","source":"fig = plt.figure(figsize=(16, 6))\nfor i in range(7):\n    fig.add_subplot(2, 4, i + 1)\n    plt.gca().title.set_text(f'mask[{i}]')\n    plt.imshow(masks[i])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:27:54.17832Z","iopub.execute_input":"2021-11-26T09:27:54.178571Z","iopub.status.idle":"2021-11-26T09:27:55.132429Z","shell.execute_reply.started":"2021-11-26T09:27:54.178537Z","shell.execute_reply":"2021-11-26T09:27:55.130692Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"distance transform can be obtained from mean of several thin masks.","metadata":{}},{"cell_type":"code","source":"energy = np.mean(np.array(masks[:5]), axis=0)\nplt.imshow(energy)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:28:45.358085Z","iopub.execute_input":"2021-11-26T09:28:45.358807Z","iopub.status.idle":"2021-11-26T09:28:45.634534Z","shell.execute_reply.started":"2021-11-26T09:28:45.358769Z","shell.execute_reply":"2021-11-26T09:28:45.633757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"we can get instance segmented mask by watershed algorithm using distance transform.\nThis is same as masks[6]","metadata":{}},{"cell_type":"code","source":"def 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    # Marker labelling\n    markers = label(energy_ths)\n\n    labels = watershed(-energy,\n                       markers,\n                       mask=msk_ths,\n                       watershed_line=line)\n\n    return labels\n\n\ndef sigmoid(x):\n    return 1 / (1 + np.exp(-x))","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T09:28:45.656418Z","iopub.execute_input":"2021-11-26T09:28:45.65672Z","iopub.status.idle":"2021-11-26T09:28:45.662544Z","shell.execute_reply.started":"2021-11-26T09:28:45.656688Z","shell.execute_reply":"2021-11-26T09:28:45.661871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"valid_mask = watershed_energy(sigmoid(masks[0]) * 255, sigmoid(energy) * 255, 0.5, 0.7, line=True)\nplt.imshow(valid_mask)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:28:45.936941Z","iopub.execute_input":"2021-11-26T09:28:45.937377Z","iopub.status.idle":"2021-11-26T09:28:46.268512Z","shell.execute_reply.started":"2021-11-26T09:28:45.937343Z","shell.execute_reply":"2021-11-26T09:28:46.267855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train","metadata":{}},{"cell_type":"markdown","source":"### Dataset and Augmentation","metadata":{}},{"cell_type":"code","source":"class CellDataset(Dataset):\n    def __init__(self, df, transforms):\n        self.df = df\n        self.base_path = config.TRAIN_PATH\n        self.transforms = transforms\n        self.image_ids = df.id.unique().tolist()\n\n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        image_path = os.path.join(self.base_path, image_id + \".png\")\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        with open(f'{config.MASK_PATH}/{image_id}.pickle', 'rb') as f:\n            masks = pickle.load(f)\n\n        augmented = self.transforms(image=image, masks=masks)\n        image = augmented['image']\n        masks = augmented['masks']\n        masks = torch.tensor(np.stack(masks))\n        return {'image': image.float(), 'masks': masks.int()}\n\n    def __len__(self):\n        return len(self.image_ids)\n\n\ndata_transforms = {\n    \"train\": A.Compose([\n        A.Resize(config.IMAGE_RESIZE[0], config.IMAGE_RESIZE[1]),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n#         A.RandomRotate90(p=0.5),\n        A.ColorJitter(p=0.5),\n        A.ChannelShuffle(p=0.25),\n        A.ToGray(p=0.25),\n        A.Normalize(mean=config.RESNET_MEAN, std=config.RESNET_STD, p=1),\n        ToTensorV2()], p=1.0),\n\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-26T09:28:52.978012Z","iopub.execute_input":"2021-11-26T09:28:52.978567Z","iopub.status.idle":"2021-11-26T09:28:52.990015Z","shell.execute_reply.started":"2021-11-26T09:28:52.978529Z","shell.execute_reply":"2021-11-26T09:28:52.989208Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Utils","metadata":{}},{"cell_type":"code","source":"def 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    intersection = intersection[1:, 1:]  # exclude background\n    union = union[1:, 1:]\n    union[union == 0] = 1e-9\n    iou = intersection / union\n\n    return iou\n\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 [n_truths x n_preds]): 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_negatives = np.sum(matches, axis=1) == 0  # Missed objects\n    false_positives = np.sum(matches, axis=0) == 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\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.numpy().astype(int), pred.numpy().astype(int)) for truth, pred in zip(truths, preds)]\n\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":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T09:28:53.947811Z","iopub.execute_input":"2021-11-26T09:28:53.948287Z","iopub.status.idle":"2021-11-26T09:28:53.972277Z","shell.execute_reply.started":"2021-11-26T09:28:53.948247Z","shell.execute_reply":"2021-11-26T09:28:53.971598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Train and valid loop","metadata":{}},{"cell_type":"code","source":"def train_loop(model, optimizer, loader, criterion):\n    losses, lrs = [], []\n    model.train()\n    optimizer.zero_grad()\n    for d in loader:\n        y = d['masks'].to(config.device)\n        pred_y = model(d['image'].to(config.device))\n        loss = criterion(pred_y, y[:, :-1].float())\n        losses.append(loss.item())\n        step_lr = np.array([param_group[\"lr\"] for param_group in optimizer.param_groups]).mean()\n        lrs.append(step_lr)\n        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n    return np.array(losses).mean(), np.array(lrs).mean()\n\n\ndef valid_loop(model, loader, criterion):\n    losses, true_masks, pred_masks, pred_energys = [], [], [], []\n    model.eval()\n    for d in loader:\n        with torch.no_grad():\n            y = d['masks'].to(config.device)\n            pred_y = model(d['image'].to(config.device))\n            loss = criterion(pred_y, y[:, :-1].float())\n        losses.append(loss.item())\n        energy = torch.mean(F.sigmoid(pred_y[:, :(config.mask_len - 1)]), dim=1)\n        pred_masks.append(F.sigmoid(pred_y[:, 0].cpu()))\n        true_masks.append(y[:, -1].cpu())\n        pred_energys.append(energy.cpu())\n    pred_masks = torch.cat(pred_masks)\n    true_masks = torch.cat(true_masks)\n    pred_energys = torch.cat(pred_energys)\n    return np.array(losses).mean(), true_masks, pred_masks, pred_energys","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-11-26T09:28:54.536846Z","iopub.execute_input":"2021-11-26T09:28:54.537396Z","iopub.status.idle":"2021-11-26T09:28:54.549782Z","shell.execute_reply.started":"2021-11-26T09:28:54.537359Z","shell.execute_reply":"2021-11-26T09:28:54.54887Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"skf = StratifiedKFold(n_splits=config.N_FOLD, shuffle=True)\nfor fold, (_, valid_idx) in enumerate(skf.split(df_train, df_train.cell_type)):\n    df_train.loc[valid_idx, 'fold'] = fold\n    \nfor fold in range(config.N_FOLD):\n    print(f\"Fold: {fold}\")\n    train_dset = CellDataset(df_train.query(f\"fold!={fold}\"), data_transforms['train'])\n    valid_dset = CellDataset(df_train.query(f\"fold=={fold}\"), data_transforms['valid'])\n\n    train_loader = DataLoader(train_dset, batch_size=config.BS,\n                              pin_memory=True, shuffle=True, num_workers=config.num_workers,\n                              worker_init_fn=lambda x: np.random.seed(torch.initial_seed() // 2 ** 32 + x))\n    valid_loader = DataLoader(valid_dset, batch_size=config.BS * 2,\n                              pin_memory=True, shuffle=False, drop_last=False, num_workers=config.num_workers)\n\n    model = smp.Unet(config.model_name, encoder_weights=\"imagenet\", activation=None, classes=config.mask_len)\n    model = model.to(config.device)\n\n    optimizer = optim.Adam(model.parameters(), lr=config.LR)\n    criterion = smp.losses.JaccardLoss(mode='multilabel')\n    scheduler = CosineAnnealingLR(optimizer, T_max=config.N_EPOCH, eta_min=config.min_LR)\n\n    valid_best_score = 0.\n    for epoch in tqdm(range(config.N_EPOCH)):\n        train_loss, lrs = train_loop(model, optimizer, train_loader, criterion)\n        valid_loss, valid_mask, valid_pred_mask, valid_pred_energy = valid_loop(model, valid_loader, criterion)\n\n        preds_wt = torch.stack([torch.tensor(watershed_energy(pred.numpy() * 255, energy.numpy() * 255, 0.5, 0.7)) for pred, energy in zip(valid_pred_mask, valid_pred_energy)])\n        valid_score_energy = iou_map(valid_mask, preds_wt)\n        if valid_score_energy > valid_best_score:\n            print(f\"epoch: {epoch}, train_loss: {train_loss:.3f}, valid_loss: {valid_loss:.3f}, meanIoU: {valid_score_energy:.3f}\")\n            valid_best_score = valid_score_energy\n            torch.save(model.state_dict(), f'{config.MODEL_PATH}/{config.model_name}_{fold}.pth')\n        scheduler.step()\n    break","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:28:55.186867Z","iopub.execute_input":"2021-11-26T09:28:55.187412Z","iopub.status.idle":"2021-11-26T09:32:13.331331Z","shell.execute_reply.started":"2021-11-26T09:28:55.187374Z","shell.execute_reply":"2021-11-26T09:32:13.328111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}