{"cells":[{"metadata":{},"cell_type":"markdown","source":"# RANZCR 1st Place Solution Training Segmentation Model\n\nHi all,\n\nWe're very exciting to writing this notebook and the summary of our solution here.\n\nOur final pipeline has 4 training stages but the minimal pipeline I show here has only 2 stages.\n\nThe 5-fold model trained with this minimal pipeline is sufficient to achieve CV 0.968-0.969 and pub/pvt LB 0.972\n\nI published 3 notebooks to demonstrate how our MINIMAL pipeline works.\n\n* Stage1: Segmentation (This notebook)\n* Stage2: Classification (https://www.kaggle.com/haqishen/ranzcr-1st-place-soluiton-cls-model-small-ver)\n* Inference (https://www.kaggle.com/haqishen/ranzcr-1st-place-soluiton-inference-small-ver)\n\nThis notebook shows how we can train a segmentation model (b1 w/ input size 1024) of our training pipeline.\n\nTo get a similar performance of a single segmentation model to ours, you only need to do:\n\n* set DEBUG to `False`\n* train for 30 epochs for each fold\n* change the model type to effnet-b5/b6/b7\n\nNote that when `DEBUG = True` the model takes roughly 10h to train a fold on a single V100, so it is **NOT** possible to train on the kaggle kernel.\nIt tooks around 2 days for training all 5 fold segmentation models.\n\nAfter 5 fold segmentation models are all trained, we'll generate masks for all images, including those w/o annotation ones.\n\nWe'll then use the predicted masks as part of input to train out stage 2 —— classification models.\n\nOur brief summary of winning solution: https://www.kaggle.com/c/ranzcr-clip-catheter-line-classification/discussion/226633\n\n\n# Main Ideas\n\n* amp\n* convert annotation to 2ch mask\n* optimized augmentation methods\n* bce loss\n* warmup + cosine scheduler\n* predict images w/ annotations as OOF, predict images w/o annotations by 5 fold ensemble\n\n\n# Thanks!"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip -q install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DEBUG = True\n\n# Add smp to kernel w/o Internet\nimport sys\nsys.path = [\n    '../input/smp20210127/segmentation_models.pytorch-master/segmentation_models.pytorch-master/',\n    '../input/smp20210127/EfficientNet-PyTorch-master/EfficientNet-PyTorch-master',\n    '../input/smp20210127/pytorch-image-models-master/pytorch-image-models-master',\n    '../input/smp20210127/pretrained-models.pytorch-master/pretrained-models.pytorch-master',\n] + sys.path","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# libraries\nimport os\nimport time\nimport subprocess\nimport numpy as np\nimport pandas as pd\nimport ast\nimport cv2\nimport PIL.Image\nimport matplotlib.pyplot as plt\n%matplotlib inline\nimport seaborn as sns\nimport torch\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms as transforms\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nfrom torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR\nfrom warmup_scheduler import GradualWarmupScheduler\nimport albumentations\nimport torch.cuda.amp as amp\nimport segmentation_models_pytorch as smp\nfrom tqdm import tqdm_notebook as tqdm\n\nscaler = amp.GradScaler()\ndevice = torch.device('cuda')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"kernel_type = 'unet++b1_2cbce_1024T15tip_lr1e4_bs4_augv2_30epo'\nenet_type = 'timm-efficientnet-b1'\ndata_dir = '../input/ranzcr-clip-catheter-line-classification'\nnum_workers = 16\nimage_size = 1024\nbatch_size = 4\ninit_lr = 1e-4\nwarmup_epo = 1\n# If DEBUG == True, only run 3 epochs per fold\ncosine_epo = 29 if not DEBUG else 2\nn_epochs = warmup_epo + cosine_epo\nuse_amp = True\nimage_folder = 'train'\n\nlog_dir = './logs'\nmodel_dir = './models'\nos.makedirs(log_dir, exist_ok=True)\nos.makedirs(model_dir, exist_ok=True)\nlog_file = os.path.join(log_dir, f'log_{kernel_type}.txt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = pd.read_csv('../input/ranzcr-fold/train_v2.csv')\n\n# If DEBUG == True, use only 100 samples to run.\ndf_train = pd.concat([\n    df_train.query('fold == 0').sample(20),\n    df_train.query('fold == 1').sample(20),\n    df_train.query('fold == 2').sample(20),\n    df_train.query('fold == 3').sample(20),\n    df_train.query('fold == 4').sample(20),\n]) if DEBUG else df_train\n\ndf_train_anno = pd.read_csv(os.path.join(data_dir, 'train_annotations.csv'))\ndf_train.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RANZCRDataset(Dataset):\n\n    def __init__(self, df, mode, transform=None):\n\n        self.df = df.reset_index(drop=True)\n        self.mode = mode\n        self.transform = transform\n\n    def __len__(self):\n        return self.df.shape[0]\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        image = cv2.imread(os.path.join(data_dir, image_folder, row.StudyInstanceUID + '.jpg'))[:, :, ::-1]\n\n        if self.mode == 'test':\n            mask = None\n            res = self.transform(image=image)\n        else:\n            df_this = df_train_anno.query(f'StudyInstanceUID == \"{row.StudyInstanceUID}\"')\n            mask = np.zeros((image.shape[0], image.shape[1], 2)).astype(np.uint8)\n            for _, anno in df_this.iterrows():\n                anno_this = np.array(ast.literal_eval(anno[\"data\"]))\n                mask1 = mask[:, :, 0].copy()\n                mask1 = cv2.polylines(mask1, np.int32([anno_this]), isClosed=False, color=1, thickness=15, lineType=16)\n                mask[:, :, 0] = mask1\n                mask2 = mask[:, :, 1].copy()\n                mask2 = cv2.circle(mask2, (anno_this[0][0], anno_this[0][1]), radius=15, color=1, thickness=25)\n                mask2 = cv2.circle(mask2, (anno_this[-1][0], anno_this[-1][1]), radius=15, color=1, thickness=25)\n                mask[:, :, 1] = mask2\n\n            mask = cv2.resize(mask ,(image_size, image_size))\n            mask = (mask > 0.5).astype(np.uint8)\n            res = self.transform(image=image, mask=mask)\n\n        image = res['image'].astype(np.float32).transpose(2, 0, 1) / 255.\n\n        if self.mode == 'test':\n            return torch.tensor(image)\n        else:\n            mask = res['mask'].astype(np.float32)\n            mask = mask.transpose(2, 0, 1).clip(0, 1)\n            return torch.tensor(image), torch.tensor(mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    albumentations.Resize(image_size, image_size),                                    \n    albumentations.HorizontalFlip(p=0.5),\n    albumentations.RandomBrightness(limit=0.1, p=0.75),\n    albumentations.ShiftScaleRotate(shift_limit=0.2, scale_limit=0.2, rotate_limit=30, border_mode=0, p=0.75),\n    albumentations.Cutout(max_h_size=int(image_size * 0.3), max_w_size=int(image_size * 0.3), num_holes=1, p=0.75),\n])\ntransforms_val = albumentations.Compose([\n    albumentations.Resize(image_size, image_size),\n])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualization"},{"metadata":{"trusted":true},"cell_type":"code","source":"df_show = df_train.query('w_anno==True').iloc[:8]\ndataset_show = RANZCRDataset(df_show, 'train', transform=transforms_train)\nfrom pylab import rcParams\nrcParams['figure.figsize'] = 20,10\n\nf, axarr = plt.subplots(1,5)\nmasks = []\nfor p in range(5):\n    img, mask = dataset_show[p]\n    img[0] = img[0]\n    axarr[p].imshow(img.transpose(0, 1).transpose(1,2))\n    masks.append(mask)\n\nf, axarr = plt.subplots(1,5)\nfor p in range(5):\n    axarr[p].imshow(masks[p][0])\n\nf, axarr = plt.subplots(1,5)\nfor p in range(5):\n    axarr[p].imshow(masks[p][1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SegModel(nn.Module):\n    def __init__(self, backbone):\n        super(SegModel, self).__init__()\n        self.seg = smp.UnetPlusPlus(encoder_name=backbone, encoder_weights='imagenet', classes=2, activation=None)\n    def forward(self,x):\n        global_features = self.seg.encoder(x)\n        seg_features = self.seg.decoder(*global_features)\n        seg_features = self.seg.segmentation_head(seg_features)\n        return seg_features\n\nm = SegModel(enet_type)\nm(torch.rand(2,3,224,224)).shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"criterion = nn.BCEWithLogitsLoss()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# LR Scheduler"},{"metadata":{"trusted":true},"cell_type":"code","source":"class GradualWarmupSchedulerV2(GradualWarmupScheduler):\n    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):\n        super(GradualWarmupSchedulerV2, self).__init__(optimizer, multiplier, total_epoch, after_scheduler)\n    def get_lr(self):\n        if self.last_epoch > self.total_epoch:\n            if self.after_scheduler:\n                if not self.finished:\n                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]\n                    self.finished = True\n                return self.after_scheduler.get_lr()\n            return [base_lr * self.multiplier for base_lr in self.base_lrs]\n        if self.multiplier == 1.0:\n            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]\n        else:\n            return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]\n\noptimizer = optim.Adam(m.parameters(), lr=init_lr)\nscheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cosine_epo)\nscheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\nlrs = []\nfor epoch in range(1, n_epochs+1):\n    scheduler_warmup.step(epoch-1)\n    lrs.append(optimizer.param_groups[0][\"lr\"])\nrcParams['figure.figsize'] = 20,3\nplt.plot(lrs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train & Valid Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(model, loader, optimizer):\n\n    model.train()\n    train_loss = []\n    bar = tqdm(loader)\n    for (data, mask) in bar:\n\n        optimizer.zero_grad()\n        data, mask = data.to(device), mask.to(device)\n\n        with amp.autocast():\n            logits = model(data)\n            loss = criterion(logits, mask)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        loss_np = loss.item()\n        train_loss.append(loss_np)\n        smooth_loss = sum(train_loss[-50:]) / min(len(train_loss), 50)\n        bar.set_description('loss: %.4f, smth: %.4f' % (loss_np, smooth_loss))\n\n    return np.mean(train_loss)\n\n\ndef valid_epoch(model, loader, get_output=False):\n\n    model.eval()\n    val_loss = []\n    LOGITS = []\n    with torch.no_grad():\n        for (data, mask) in tqdm(loader):\n            data, mask = data.to(device), mask.to(device)\n            logits = model(data)\n            loss = criterion(logits, mask)\n            val_loss.append(loss.item())\n            LOGITS.append(logits.cpu())\n\n    if get_output:\n        LOGITS = torch.cat(LOGITS, 0).float().sigmoid()\n        return LOGITS\n    else:\n        val_loss = np.mean(val_loss)\n        return val_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Run!"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(fold):\n    content = 'Fold: ' + str(fold)\n    print(content)\n    with open(log_file, 'a') as appender:\n        appender.write(content + '\\n')\n    train_ = df_train.query(f'w_anno==True and fold!={fold}').copy()\n    valid_ = df_train.query(f'w_anno==True and fold=={fold}').copy()\n\n    dataset_train = RANZCRDataset(train_, 'train', transform=transforms_train)\n    dataset_valid = RANZCRDataset(valid_, 'valid', transform=transforms_val)\n    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)\n    valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n\n    model = SegModel(enet_type)\n    model = model.to(device)\n    val_loss_min = np.Inf\n    model_file = os.path.join(model_dir, f'{kernel_type}_best_fold{fold}.pth')\n\n    optimizer = optim.Adam(model.parameters(), lr=init_lr)\n    scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cosine_epo)\n    scheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n    for epoch in range(1, n_epochs+1):\n        print(time.ctime(), 'Epoch:', epoch)\n        scheduler_warmup.step(epoch-1)\n\n        train_loss = train_epoch(model, train_loader, optimizer)\n        val_loss = valid_epoch(model, valid_loader)\n\n        content = time.ctime() + ' ' + f'Fold {fold} Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, train loss: {train_loss:.5f}, valid loss: {(val_loss):.5f}.'\n        print(content)\n        with open(log_file, 'a') as appender:\n            appender.write(content + '\\n')\n\n        if val_loss_min > val_loss:\n            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(val_loss_min, val_loss))\n            torch.save(model.state_dict(), model_file)\n            val_loss_min = val_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold in range(5):\n    run(fold)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Generate Mask"},{"metadata":{"trusted":true},"cell_type":"code","source":"output_dir = f'mask_{kernel_type}'\nos.makedirs(output_dir, exist_ok=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Part 1, generate mask for those images with annotations. To prevent leaks, use only the model corresponding to the fold\nfor fold in range(5):\n    valid_ = df_train.query(f'w_anno==True and fold=={fold}').copy()\n    dataset_valid = RANZCRDataset(valid_, 'valid', transform=transforms_val)\n    valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n\n    model = SegModel(enet_type)\n    model = model.to(device)\n    model_file = os.path.join(model_dir, f'{kernel_type}_best_fold{fold}.pth')\n    model.load_state_dict(torch.load(model_file), strict=False)\n    model.eval()\n    \n    outputs = valid_epoch(model, valid_loader, get_output=True).numpy()\n\n    for i, (_, row) in tqdm(enumerate(valid_.iterrows())):\n        png = (outputs[i] * 255).astype(np.uint8).transpose(1,2,0)\n        # add a channel to make it able to be saved as .png\n        png = np.concatenate([png, np.zeros((png.shape[0], png.shape[1], 1))], -1)\n        cv2.imwrite(os.path.join(output_dir, f'{row.StudyInstanceUID}.png'), png)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Part 2, For those images without annotations, use 5-fold models to predict and take the mean value\ndf_train_wo_anno = df_train.query(f'w_anno==False').copy().reset_index(drop=True)\ndataset_test = RANZCRDataset(df_train_wo_anno, 'test', transform=transforms_val)\ntest_loader = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n\nmodels = []\nfor fold in range(5):\n    model = SegModel(enet_type)\n    model = model.to(device)\n    model_file = os.path.join(model_dir, f'{kernel_type}_best_fold{fold}.pth')\n    model.load_state_dict(torch.load(model_file), strict=False)\n    model.eval()\n    models.append(model)\n\nwith torch.no_grad():\n    for batch_id, data in tqdm(enumerate(test_loader), total=len(test_loader)):\n        data = data.to(device)\n        outputs = torch.stack([model(data).sigmoid() for model in models], 0).mean(0).cpu().numpy()\n        for i in range(outputs.shape[0]):\n            row = df_train_wo_anno.loc[batch_id*batch_size + i]\n            png = (outputs[i] * 255).astype(np.uint8).transpose(1,2,0)\n            # add a channel to make it able to be saved as .png\n            png = np.concatenate([png, np.zeros((png.shape[0], png.shape[1], 1))], -1)\n            cv2.imwrite(os.path.join(output_dir, f'{row.StudyInstanceUID}.png'), png)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!zip -r -q pred_masks.zip {output_dir}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -rf {output_dir}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ls","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Appendix: My Local Training Log Fold 0\n\nYou should get similar log to mine only by changing `DEBUG = False` and train on your local machine with >= 16GB GPU.\n\n\n```\nWed Mar 17 01:08:11 2021 Fold 0 Epoch 1, lr: 0.0001000, train loss: 0.08337, valid loss: 0.01769.\nWed Mar 17 01:28:57 2021 Fold 0 Epoch 2, lr: 0.0010000, train loss: 0.01330, valid loss: 0.01110.\nWed Mar 17 01:49:42 2021 Fold 0 Epoch 3, lr: 0.0010000, train loss: 0.01107, valid loss: 0.01013.\nWed Mar 17 02:10:28 2021 Fold 0 Epoch 4, lr: 0.0009883, train loss: 0.01042, valid loss: 0.00958.\nWed Mar 17 02:31:13 2021 Fold 0 Epoch 5, lr: 0.0009738, train loss: 0.01002, valid loss: 0.00989.\nWed Mar 17 02:51:57 2021 Fold 0 Epoch 6, lr: 0.0009538, train loss: 0.00971, valid loss: 0.00918.\nWed Mar 17 03:12:41 2021 Fold 0 Epoch 7, lr: 0.0009284, train loss: 0.00949, valid loss: 0.00925.\nWed Mar 17 03:33:25 2021 Fold 0 Epoch 8, lr: 0.0008980, train loss: 0.00924, valid loss: 0.00898.\nWed Mar 17 03:54:11 2021 Fold 0 Epoch 9, lr: 0.0008630, train loss: 0.00912, valid loss: 0.00887.\nWed Mar 17 04:14:57 2021 Fold 0 Epoch 10, lr: 0.0008237, train loss: 0.00890, valid loss: 0.00857.\nWed Mar 17 04:35:42 2021 Fold 0 Epoch 11, lr: 0.0007806, train loss: 0.00879, valid loss: 0.00856.\nWed Mar 17 04:56:25 2021 Fold 0 Epoch 12, lr: 0.0007342, train loss: 0.00871, valid loss: 0.00839.\nWed Mar 17 05:17:10 2021 Fold 0 Epoch 13, lr: 0.0006851, train loss: 0.00853, valid loss: 0.00836.\nWed Mar 17 05:37:53 2021 Fold 0 Epoch 14, lr: 0.0006338, train loss: 0.00845, valid loss: 0.00840.\nWed Mar 17 05:58:37 2021 Fold 0 Epoch 15, lr: 0.0005809, train loss: 0.00827, valid loss: 0.00821.\nWed Mar 17 06:19:23 2021 Fold 0 Epoch 16, lr: 0.0005271, train loss: 0.00821, valid loss: 0.00814.\nWed Mar 17 06:40:07 2021 Fold 0 Epoch 17, lr: 0.0004729, train loss: 0.00813, valid loss: 0.00824.\nWed Mar 17 07:00:51 2021 Fold 0 Epoch 18, lr: 0.0004191, train loss: 0.00809, valid loss: 0.00808.\nWed Mar 17 07:21:37 2021 Fold 0 Epoch 19, lr: 0.0003662, train loss: 0.00794, valid loss: 0.00801.\nWed Mar 17 07:42:21 2021 Fold 0 Epoch 20, lr: 0.0003149, train loss: 0.00783, valid loss: 0.00798.\nWed Mar 17 08:03:05 2021 Fold 0 Epoch 21, lr: 0.0002658, train loss: 0.00782, valid loss: 0.00797.\nWed Mar 17 08:23:52 2021 Fold 0 Epoch 22, lr: 0.0002194, train loss: 0.00768, valid loss: 0.00790.\nWed Mar 17 08:44:36 2021 Fold 0 Epoch 23, lr: 0.0001763, train loss: 0.00765, valid loss: 0.00794.\nWed Mar 17 09:05:20 2021 Fold 0 Epoch 24, lr: 0.0001370, train loss: 0.00756, valid loss: 0.00788.\nWed Mar 17 09:26:06 2021 Fold 0 Epoch 25, lr: 0.0001020, train loss: 0.00758, valid loss: 0.00789.\nWed Mar 17 09:46:58 2021 Fold 0 Epoch 26, lr: 0.0000716, train loss: 0.00750, valid loss: 0.00784.\nWed Mar 17 10:07:39 2021 Fold 0 Epoch 27, lr: 0.0000462, train loss: 0.00745, valid loss: 0.00786.\nWed Mar 17 10:28:26 2021 Fold 0 Epoch 28, lr: 0.0000262, train loss: 0.00743, valid loss: 0.00786.\nWed Mar 17 10:49:13 2021 Fold 0 Epoch 29, lr: 0.0000117, train loss: 0.00746, valid loss: 0.00784.\nWed Mar 17 11:09:55 2021 Fold 0 Epoch 30, lr: 0.0000029, train loss: 0.00741, valid loss: 0.00785.\n```"},{"metadata":{},"cell_type":"markdown","source":""}],"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":4,"nbformat_minor":4}