{"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-12-10T09:36:13.137384Z","iopub.execute_input":"2021-12-10T09:36:13.137725Z","iopub.status.idle":"2021-12-10T09:36:47.981893Z","shell.execute_reply.started":"2021-12-10T09:36:13.137625Z","shell.execute_reply":"2021-12-10T09:36:47.980893Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm.notebook import tqdm 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\nfrom statistics import mean, stdev\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/cell-masks/train_masks\"\n    RESNET_MEAN = [0.485]\n    RESNET_STD = [0.229]\n    IMAGE_RESIZE = [512, 704]\n    LR = 5e-4\n    min_LR = 5e-5\n    device = 'cuda'\n    BS = 4\n    num_workers = 2\n    N_EPOCH = 50\n    N_FOLD = 5\n    \nos.makedirs(config.MODEL_PATH, exist_ok=True)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-10T09:36:47.984468Z","iopub.execute_input":"2021-12-10T09:36:47.984803Z","iopub.status.idle":"2021-12-10T09:36:57.283500Z","shell.execute_reply.started":"2021-12-10T09:36:47.984756Z","shell.execute_reply":"2021-12-10T09:36:57.282754Z"},"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-12-10T09:36:57.285000Z","iopub.execute_input":"2021-12-10T09:36:57.285248Z","iopub.status.idle":"2021-12-10T09:36:57.888967Z","shell.execute_reply.started":"2021-12-10T09:36:57.285214Z","shell.execute_reply":"2021-12-10T09:36:57.887973Z"},"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}/mask_{image_id}.pkl', 'rb') as f:\n    masks = pickle.load(f)\n\nmasks.shape","metadata":{"execution":{"iopub.status.busy":"2021-12-10T09:36:57.890274Z","iopub.execute_input":"2021-12-10T09:36:57.890601Z","iopub.status.idle":"2021-12-10T09:36:57.928865Z","shell.execute_reply.started":"2021-12-10T09:36:57.890560Z","shell.execute_reply":"2021-12-10T09:36:57.927963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train","metadata":{}},{"cell_type":"markdown","source":"### Dataset and Augmentation","metadata":{}},{"cell_type":"code","source":"from torchvision import transforms as T\nfrom PIL import Image\n\nclass 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 = Image.open(image_path)\n        \n        with open(f'{config.MASK_PATH}/mask_{image_id}.pkl', 'rb') as f:\n            masks = pickle.load(f)\n\n        masks = Image.fromarray(masks)\n        image = self.transforms(image)\n        masks = masks.resize(config.IMAGE_RESIZE)\n        masks = T.ToTensor()(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\": T.Compose([\n        T.Resize(config.IMAGE_RESIZE),\n#         T.RandomHorizontalFlip(p=0.5),\n#         T.RandomVerticalFlip(p=0.5),\n        T.ToTensor(),\n        T.Normalize(mean=config.RESNET_MEAN, std=config.RESNET_STD),\n    ]),\n\n    \"valid\": T.Compose([\n        T.Resize(config.IMAGE_RESIZE),\n        T.ToTensor(),\n        T.Normalize(mean=config.RESNET_MEAN, std=config.RESNET_STD),\n    ])\n}","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-10T09:36:57.931304Z","iopub.execute_input":"2021-12-10T09:36:57.931568Z","iopub.status.idle":"2021-12-10T09:36:57.942325Z","shell.execute_reply.started":"2021-12-10T09:36:57.931532Z","shell.execute_reply":"2021-12-10T09:36:57.941206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Utils","metadata":{}},{"cell_type":"code","source":"def get_threshold(Y, pred):\n    scores = list(pred.ravel())\n    mask = list(Y.ravel())\n    \n    idxs=np.argsort(scores)[::-1]\n    mask_sorted=np.array(mask)[idxs]\n    sum_mask_one=np.cumsum(mask_sorted)\n    IoU=sum_mask_one/(np.arange(1,len(mask_sorted)+1)+np.sum(mask_sorted)-sum_mask_one)\n    best_IoU_idx=IoU.argmax()\n    best_threshold=scores[idxs[best_IoU_idx]]\n    best_IoU=IoU[best_IoU_idx]\n\n    return best_threshold, best_IoU","metadata":{"execution":{"iopub.status.busy":"2021-12-10T09:36:57.943810Z","iopub.execute_input":"2021-12-10T09:36:57.944066Z","iopub.status.idle":"2021-12-10T09:36:57.953343Z","shell.execute_reply.started":"2021-12-10T09:36:57.944028Z","shell.execute_reply":"2021-12-10T09:36:57.952187Z"},"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.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.float())\n        losses.append(loss.item())\n        pred_masks.append(F.sigmoid(pred_y.cpu()))\n        true_masks.append(y.cpu())\n    pred_masks = torch.cat(pred_masks)\n    true_masks = torch.cat(true_masks)\n    return np.array(losses).mean(), true_masks, pred_masks","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-10T09:43:00.146257Z","iopub.execute_input":"2021-12-10T09:43:00.146941Z","iopub.status.idle":"2021-12-10T09:43:00.158222Z","shell.execute_reply.started":"2021-12-10T09:43:00.146900Z","shell.execute_reply":"2021-12-10T09:43:00.157067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Network","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\n\nclass Attention_block(nn.Module):\n    def __init__(self,F_g,F_l,F_int):\n        super(Attention_block,self).__init__()\n        self.W_g = nn.Sequential(\n            nn.Conv2d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(F_int)\n            )\n        \n        self.W_x = nn.Sequential(\n            nn.Conv2d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(F_int)\n        )\n\n        self.psi = nn.Sequential(\n            nn.Conv2d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(1),\n            nn.Sigmoid()\n        )\n        \n        self.relu = nn.ReLU(inplace=True)\n        \n    def forward(self,g,x):\n        g1 = self.W_g(g)\n        x1 = self.W_x(x)\n        psi = self.relu(g1+x1)\n        psi = self.psi(psi)\n\n        return x*psi","metadata":{"execution":{"iopub.status.busy":"2021-12-10T09:36:57.968641Z","iopub.execute_input":"2021-12-10T09:36:57.969137Z","iopub.status.idle":"2021-12-10T09:36:57.979628Z","shell.execute_reply.started":"2021-12-10T09:36:57.969099Z","shell.execute_reply":"2021-12-10T09:36:57.978813Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class conv_block(nn.Module):\n    def __init__(self,ch_in,ch_out):\n        super(conv_block,self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),\n            nn.BatchNorm2d(ch_out),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(ch_out, ch_out, kernel_size=3,stride=1,padding=1,bias=True),\n            nn.BatchNorm2d(ch_out),\n            nn.ReLU(inplace=True)\n        )\n\n\n    def forward(self,x):\n        x = self.conv(x)\n        return x\n\nclass up_conv(nn.Module):\n    def __init__(self,ch_in,ch_out):\n        super(up_conv,self).__init__()\n        self.up = nn.Sequential(\n            nn.Upsample(scale_factor=2),\n            nn.Conv2d(ch_in,ch_out,kernel_size=3,stride=1,padding=1,bias=True),\n            nn.BatchNorm2d(ch_out),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self,x):\n        x = self.up(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2021-12-10T09:36:57.981042Z","iopub.execute_input":"2021-12-10T09:36:57.981316Z","iopub.status.idle":"2021-12-10T09:36:57.990981Z","shell.execute_reply.started":"2021-12-10T09:36:57.981279Z","shell.execute_reply":"2021-12-10T09:36:57.990104Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AttU_Net(nn.Module):\n    def __init__(self,img_ch=3,output_ch=1,scale_factor=1):\n        super(AttU_Net,self).__init__()\n        \n        self.Maxpool = nn.MaxPool2d(kernel_size=2,stride=2)\n\n        filters = [64, 128, 256, 512, 1024]\n        filters = [int(x / scale_factor) for x in filters]\n\n        self.Conv1 = conv_block(ch_in=img_ch,ch_out=filters[0])\n        self.Conv2 = conv_block(ch_in=filters[0],ch_out=filters[1])\n        self.Conv3 = conv_block(ch_in=filters[1],ch_out=filters[2])\n        self.Conv4 = conv_block(ch_in=filters[2],ch_out=filters[3])\n        self.Conv5 = conv_block(ch_in=filters[3],ch_out=filters[4])\n\n        self.Up5 = up_conv(ch_in=filters[4],ch_out=filters[3])\n        self.Att5 = Attention_block(F_g=filters[3],F_l=filters[3],F_int=filters[2])\n        self.Up_conv5 = conv_block(ch_in=filters[4], ch_out=filters[3])\n\n        self.Up4 = up_conv(ch_in=filters[3], ch_out=filters[2])\n        self.Att4 = Attention_block(F_g=filters[2],F_l=filters[2],F_int=filters[1])\n        self.Up_conv4 = conv_block(ch_in=filters[3], ch_out=filters[2])\n        \n        self.Up3 = up_conv(ch_in=filters[2], ch_out=filters[1])\n        self.Att3 = Attention_block(F_g=filters[1],F_l=filters[1],F_int=filters[0])\n        self.Up_conv3 = conv_block(ch_in=filters[2], ch_out=filters[1])\n        \n        self.Up2 = up_conv(ch_in=filters[1], ch_out=filters[0])\n        self.Att2 = Attention_block(F_g=filters[0],F_l=filters[0],F_int=filters[0] // 2)\n        self.Up_conv2 = conv_block(ch_in=filters[1], ch_out=filters[0])\n\n        self.Conv_1x1 = nn.Conv2d(filters[0],output_ch,kernel_size=1,stride=1,padding=0)\n\n\n    def forward(self,x):\n        # encoding path\n        x1 = self.Conv1(x)\n\n        x2 = self.Maxpool(x1)\n        x2 = self.Conv2(x2)\n        \n        x3 = self.Maxpool(x2)\n        x3 = self.Conv3(x3)\n\n        x4 = self.Maxpool(x3)\n        x4 = self.Conv4(x4)\n\n        x5 = self.Maxpool(x4)\n        x5 = self.Conv5(x5)\n\n        # decoding + concat path\n        d5 = self.Up5(x5)\n        x4 = self.Att5(g=d5,x=x4)\n        d5 = torch.cat((x4,d5),dim=1)        \n        d5 = self.Up_conv5(d5)\n        \n        d4 = self.Up4(d5)\n        x3 = self.Att4(g=d4,x=x3)\n        d4 = torch.cat((x3,d4),dim=1)\n        d4 = self.Up_conv4(d4)\n\n        d3 = self.Up3(d4)\n        x2 = self.Att3(g=d3,x=x2)\n        d3 = torch.cat((x2,d3),dim=1)\n        d3 = self.Up_conv3(d3)\n\n        d2 = self.Up2(d3)\n        x1 = self.Att2(g=d2,x=x1)\n        d2 = torch.cat((x1,d2),dim=1)\n        d2 = self.Up_conv2(d2)\n\n        d1 = self.Conv_1x1(d2)\n\n        return d1","metadata":{"execution":{"iopub.status.busy":"2021-12-10T09:36:57.992295Z","iopub.execute_input":"2021-12-10T09:36:57.992939Z","iopub.status.idle":"2021-12-10T09:36:58.014533Z","shell.execute_reply.started":"2021-12-10T09:36:57.992903Z","shell.execute_reply":"2021-12-10T09:36:58.013498Z"},"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 = AttU_Net(img_ch=1,output_ch=1,scale_factor=1)\n    model = model.to(config.device)\n\n    optimizer = optim.Adam(model.parameters(), lr=config.LR)\n    criterion = smp.losses.JaccardLoss(mode='binary')\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        img_thresholds = []         # one for each image\n        img_IoUs = []\n        train_loss, lrs = train_loop(model, optimizer, train_loader, criterion)\n        valid_loss, valid_mask, valid_pred_mask = valid_loop(model, valid_loader, criterion)\n#         for i in range(valid_mask.shape[0]):\n#             best_img_threshold, best_img_IoU = get_threshold(valid_mask[i], valid_pred_mask[i])\n#             img_thresholds.append(best_img_threshold)\n#             img_IoUs.append(best_img_IoU)\n#         best_threshold = np.mean(img_thresholds)\n#         best_threshold_spread = np.std(img_thresholds)\n#         avg_IoU = mean(img_IoUs)\n#         if avg_IoU > valid_best_score:\n        print(f\"epoch: {epoch}, train_loss: {train_loss:.3f}, valid_loss: {valid_loss:.3f}\")\n#             print(f\"Best threshold: {best_threshold:.3g} (+-{best_threshold_spread:.3g}), Avg. Train IoU: {avg_IoU:.3f}\")\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-12-10T09:43:08.898136Z","iopub.execute_input":"2021-12-10T09:43:08.898739Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}