{"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":"<center><h2 style='color:red'>HuBMAP | Pytorch | ViT for Segmentation [Train]<hr>By Kassem@elcaiseri</h2></center>","metadata":{}},{"cell_type":"markdown","source":"This Code depends on: <br>\n- The lossses are from **https://github.com/JunMa11/SegLoss**\n- The base codes are from here: **https://www.kaggle.com/vineeth1999/hubmap-pytorch-efficientunet-offline**\n\n**Update:**\n- It's the official code. https://github.com/Beckschen/TransUNet\n\n\n- It's the paper of the TransUnet. https://arxiv.org/pdf/2102.04306v1.pdf (Thanks for @luciusk)\n","metadata":{}},{"cell_type":"markdown","source":"## Introduction <br>\nImplementation of **self attention mechanisms** for computer vision in PyTorch with einsum and einops. Focused on **segmentation**  self-attention modules.\n\n### Note: This is a baseline code","metadata":{}},{"cell_type":"markdown","source":"**Update**\n- This code starts from vision 10\n- **V15**: - Use DiceBCELoss - increase ViT blocks - \n- **V16**: - Apply Official code.\n          - Add more options for loss function.","metadata":{}},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import GroupKFold\n\nimport torch\nimport torchvision\nfrom torch import nn\nfrom torch.optim import Adam\nfrom torchvision import transforms\nfrom torch.nn import functional as F\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau, CosineAnnealingWarmRestarts\n\nimport cv2\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\n\nfrom scipy.ndimage.interpolation import zoom\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport os\nimport sys\nimport time\nimport random","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Hyperparameters","metadata":{}},{"cell_type":"code","source":"class CFG:\n    data = 256 #512\n    debug=False\n    apex=False\n    print_freq=100\n    num_workers=4\n    img_size=256 # appropriate input size for encoder \n    scheduler='CosineAnnealingWarmRestarts' # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts']\n    epoch=5 # Change epochs\n    criterion= 'Lovasz' #'DiceBCELoss' # ['DiceLoss', 'Hausdorff', 'Lovasz']\n    base_model='Unet' # ['Unet']\n    encoder = 'vit' # ['attention','efficientnet-b5'] or other encoders from smp\n    lr=1e-4\n    min_lr=1e-6\n    batch_size=4\n    weight_decay=1e-6\n    gradient_accumulation_steps=1\n    seed=2021\n    n_fold=5\n    trn_fold= 0 #[0, 1, 2, 3, 4]\n    train=True\n    inference=False\n    optimizer = 'Adam'\n    T_0=10\n    # N=5 \n    # M=9\n    T_max=10\n    #factor=0.2\n    #patience=4\n    #eps=1e-6\n    smoothing=1\n    in_channels=3\n    vit_blocks=12 #[8, 12]\n    vit_linear=1024 #1024\n    classes=1\n    MODEL_NAME = 'R50-ViT-B_16'\n\n\nmain_dir = '../input/hubmap-256x256'\ntrain_dir = '../input/hubmap-256x256/train'\nmasks_dir = '../input/hubmap-256x256/masks'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from @Iafoss comment\ndef seed_torch(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    \n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    #the following line gives ~10% speedup\n    #but may lead to some stochasticity in the results \n    torch.backends.cudnn.benchmark = True # from @Iafoss comment\n\nseed_torch(seed=CFG.seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv('../input/hubmap-kidney-segmentation/train.csv')\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_transform(mode='base'):\n    if mode == 'base':\n        base_transform = A.Compose([\n            A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n            A.HorizontalFlip(),\n            A.VerticalFlip(),\n            A.RandomRotate90(),\n            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=20, p=0.4, \n                             border_mode=cv2.BORDER_REFLECT),\n            A.OneOf([\n                A.OpticalDistortion(p=0.4),\n                A.GridDistortion(p=.1),\n                A.IAAPiecewiseAffine(p=0.4),\n            ], p=0.3),\n            A.OneOf([\n                A.HueSaturationValue(10,15,10),\n                A.CLAHE(clip_limit=3),\n                A.RandomBrightnessContrast(),            \n            ], p=0.4),\n            ToTensorV2()\n        ], p=1.0)\n        return base_transform\n    \n    elif mode == 'rand':\n        rand_transform = A.Compose([\n                RandAugment(CFG.N, CFG.M),\n                A.Transpose(p=0.5),\n                A.VerticalFlip(p=0.5),\n                A.HorizontalFlip(p=0.5),\n                A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=0, p=0.85),\n                A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n                A.Normalize(),\n                ToTensorV2()\n            ])\n        return rand_transform\n    \n    elif mode == 'strong':\n        strong_transform = A.Compose([\n                A.Transpose(p=0.5),\n                A.VerticalFlip(p=0.5),\n                A.HorizontalFlip(p=0.5),\n                A.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03, p=0.5),\n                A.OneOf([\n                        A.RandomGamma(),\n                        A.GaussNoise()           \n                    ], p=0.5),\n                A.OneOf([\n                        A.OpticalDistortion(p=0.4),\n                        A.GridDistortion(p=0.2),\n                        A.IAAPiecewiseAffine(p=0.4),\n                    ], p=0.5),\n                A.OneOf([\n                        A.HueSaturationValue(10,15,10),\n                        A.CLAHE(clip_limit=4),\n                        A.RandomBrightnessContrast(),            \n                    ], p=0.5),\n\n                A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=0, p=0.85),\n                A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n                ToTensorV2()\n            ])\n        return strong_transform\n    \n    elif mode == 'weak':\n        weak_transform = A.Compose([\n                A.Resize(CFG.img_size, CFG.img_size, p=0.5),\n                A.HorizontalFlip(),\n                A.VerticalFlip(),\n                A.RandomRotate90(),\n                A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.4, \n                                 border_mode=cv2.BORDER_REFLECT),\n                ToTensorV2()\n            ], p=1.0)\n        return weak_transform\n    \n    elif mode == 'valid':\n        val_transform = A.Compose([\n                A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n                ToTensorV2()\n            ], p=1.0)\n        return val_transform\n    \n    else:\n        print('Unknown mode.')\n    \n\nmean = np.array([0.65459856,0.48386562,0.69428385])\nstd = np.array([0.15167958,0.23584107,0.13146145])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HuBMAPDataset(Dataset):\n    def __init__(self, main_dir, df, train=True, transform=None):\n        \n        self.ids = df.id.values\n        self.fnames = [fname for fname in os.listdir(train_dir) if fname.split('_')[0] in self.ids]\n\n        self.main_dir = main_dir\n        self.df = df\n        self.train = train\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.fnames)\n    \n    def __getitem__(self, idx):\n        fname = self.fnames[idx]\n        \n        img = cv2.cvtColor(cv2.imread(os.path.join(main_dir, 'train', fname)), cv2.COLOR_BGR2RGB)\n        mask = cv2.imread(os.path.join(main_dir, 'masks', fname),cv2.IMREAD_GRAYSCALE)\n        \n        if self.transform is not None:\n            aug = self.transform(image=img, mask=mask)            \n            img, mask = aug['image'], aug['mask']\n                    \n        img = img.type('torch.FloatTensor')\n        img = img/255\n        mask = mask.type('torch.FloatTensor')\n\n        return img, mask\n          ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize augmented data","metadata":{}},{"cell_type":"code","source":"def vis_aug_data(dataset, length=6):\n    plt.figure(figsize=(15,10))\n    N = length // 2\n    for i in range(length):\n        image, mask = train_dataset[i]\n        plt.subplot(3,4,2*i+1)\n        plt.imshow(np.transpose((image), (1,2,0)))\n        plt.axis('off')\n        plt.subplot(3,4,2*i+2)\n        plt.imshow(mask)\n        plt.axis('off')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# base augmentation\ntrain_dataset = HuBMAPDataset(main_dir, train_df, train=True, transform=get_transform('base'))\n\nvis_aug_data(train_dataset, 6)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# weak augmentation\ntrain_dataset = HuBMAPDataset(main_dir, train_df, train=True, transform=get_transform('weak'))\n\nvis_aug_data(train_dataset, 6)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# strong augmentation\ntrain_dataset = HuBMAPDataset(main_dir, train_df, train=True, transform=get_transform('strong'))\n\nvis_aug_data(train_dataset, 6)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# valid augmentation\ntrain_dataset = HuBMAPDataset(main_dir, train_df, train=True, transform=get_transform('valid'))\n\nvis_aug_data(train_dataset, 6)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## DataLoader","metadata":{}},{"cell_type":"code","source":"if CFG.data==512:\n    directory_list = os.listdir('../input/hubmap-512x512/train')\nelif CFG.data==256:\n    directory_list = os.listdir('../input/hubmap-256x256/train')\ndirectory_list = [fnames.split('_')[0] for fnames in directory_list]\ndir_df = pd.DataFrame(directory_list, columns=['id'])\ndir_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"!wget https://storage.googleapis.com/vit_models/imagenet21k/R50%2BViT-B_16.npz\n#! pip install self-attention-cv\n!pip install einops\n!pip install ml_collections\ntu_path = '../input/transunet/TransUNet-main'\nsys.path.append(tu_path)","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from networks.vit_seg_modeling import VisionTransformer as ViT_seg\nfrom networks.vit_seg_modeling import CONFIGS as CONFIGS_ViT_seg\n\nconfig_vit = CONFIGS_ViT_seg[CFG.MODEL_NAME]\nconfig_vit.n_classes = 1\nconfig_vit.n_skip = 3\nconfig_vit.pretrained_path = './R50+ViT-B_16.npz'\nconfig_vit.transformer.dropout_rate = 0.2\nconfig_vit.transformer.mlp_dim = 3072\nconfig_vit.transformer.num_heads = 4\nconfig_vit.transformer.num_layers = 8\n\nconfig_vit","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ViTHuBMAP(nn.Module):\n    def __init__(self, configs=config_vit):\n        super(ViTHuBMAP, self).__init__()\n        \n        self.model = ViT_seg(configs, img_size=CFG.img_size, num_classes=CFG.classes)\n        self.model.load_from(weights=np.load(configs.pretrained_path))\n\n    \n    def forward(self, x):\n        img_segs = self.model(x)\n        \n        return img_segs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss Function","metadata":{}},{"cell_type":"code","source":"sys.path.append('../input/segloss/Segmentationloss/SegLoss-master')\nfrom losses_pytorch.hausdorff import HausdorffDTLoss\nfrom losses_pytorch.lovasz_loss import LovaszSoftmax\nfrom losses_pytorch.focal_loss import FocalLoss","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=CFG.smoothing):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        \n        return dice\n    \n    \n    \nclass DiceBCELoss(nn.Module):\n    # Formula Given above.\n    def __init__(self, weight=None, size_average=True):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=CFG.smoothing):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).mean()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(inputs.mean() + targets.mean() + smooth)  \n        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')\n        Dice_BCE = BCE + dice_loss\n        \n        return Dice_BCE.mean()\n    \n    \nclass Hausdorff_loss(nn.Module):\n    def __init__(self):\n        super(Hausdorff_loss, self).__init__()\n        \n    def forward(self, inputs, targets):\n        return HausdorffDTLoss()(inputs, targets)\n    \nclass FocalDLoss(nn.Module):\n    def __init__(self):\n        super(FocalDLoss, self).__init__()\n        \n    def forward(self, inputs, targets):\n        return FocalLoss()(inputs, targets)\n    \n    \nclass Lovasz_loss(nn.Module):\n    def __init__(self):\n        super(Lovasz_loss, self).__init__()\n        \n    def forward(self, inputs, targets):\n        return LovaszSoftmax()(inputs, targets)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.criterion == 'DiceBCELoss':\n    criterion = DiceBCELoss()\nelif CFG.criterion == 'DiceLoss':\n    criterion = DiceLoss()\nelif CFG.criterion == 'FocalLoss':\n    criterion = FocalDLoss()\nelif CFG.criterion == 'Hausdorff':\n    criterion = Hausdorff_loss()\nelif CFG.criterion == 'Lovasz':\n    criterion = Lovasz_loss()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train Function","metadata":{}},{"cell_type":"code","source":"def HuBMAPLoss(images, targets, model, device, loss_func=criterion):\n    model.to(device)\n    images = images.to(device)\n    targets = targets.to(device)\n    outputs = model(images)\n    loss_func = loss_func\n    loss = loss_func(outputs, targets)\n    return loss, outputs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_one_epoch(epoch, model, device, optimizer, scheduler, trainloader):\n    model.train()\n    t = time.time()\n    total_loss = 0\n    \n    for step, (images, targets) in enumerate(trainloader):\n        loss, outputs = HuBMAPLoss(images, targets, model, device)\n        loss.backward()\n        if ((step+1)%4==0 or (step+1)==len(trainloader)):\n            optimizer.step()\n            scheduler.step()\n            optimizer.zero_grad()\n        loss = loss.detach().item()\n        total_loss += loss\n        \n        if ((step+1)%10==0 or (step+1)==len(trainloader)):\n            print(\n                    f'epoch {epoch} train step {step+1}/{len(trainloader)}, ' + \\\n                    f'loss: {total_loss/len(trainloader):.4f}, ' + \\\n                    f'time: {(time.time() - t):.4f}', end= '\\r' if (step + 1) != len(trainloader) else '\\n'\n                )\n\n            \n        \ndef valid_one_epoch(epoch, model, device, optimizer, scheduler, validloader):\n    model.eval()\n    t = time.time()\n    total_loss = 0\n    \n    for step, (images, targets) in enumerate(validloader):\n        loss, outputs = HuBMAPLoss(images, targets, model, device)\n        loss = loss.detach().item()\n        total_loss += loss\n        \n        if ((step+1)%4==0 or (step+1)==len(validloader)):\n            scheduler.step(total_loss/len(validloader))\n        \n        if ((step+1)%10==0 or (step+1)==len(validloader)):\n            print(\n                    f'**epoch {epoch} trainz step {step+1}/{len(validloader)}, ' + \\\n                    f'loss: {total_loss/len(validloader):.4f}, ' + \\\n                    f'time: {(time.time() - t):.4f}', end= '\\r' if (step + 1) != len(validloader) else '\\n'\n                )\n\n        ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating Folds Column","metadata":{}},{"cell_type":"code","source":"FOLDS = CFG.n_fold\ngkf = GroupKFold(FOLDS)\ndir_df['Folds'] = 0\nfor fold, (tr_idx, val_idx) in enumerate(gkf.split(dir_df, groups=dir_df[dir_df.columns[0]].values)):\n    dir_df.loc[val_idx, 'Folds'] = fold\n    \ndir_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The Real Training","metadata":{}},{"cell_type":"code","source":"def prepare_train_valid_dataloader(df, fold):\n    train_ids = df[~df.Folds.isin(fold)]\n    val_ids = df[df.Folds.isin(fold)]\n    \n    train_ds = HuBMAPDataset(main_dir, train_ids, train=True, transform=get_transform('base'))\n    val_ds = HuBMAPDataset(main_dir, val_ids, train=True, transform=get_transform('valid'))\n    \n    train_loader = DataLoader(train_ds, batch_size=CFG.batch_size, pin_memory=True, shuffle=True, num_workers=CFG.num_workers)\n    val_loader = DataLoader(val_ds, batch_size=CFG.batch_size, pin_memory=True, shuffle=False, num_workers=CFG.num_workers)\n    \n    return train_loader, val_loader","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nmodel = ViTHuBMAP().to(device)\noptimizer = Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay, amsgrad=False)\n\n# scheduler setting\nif CFG.scheduler == 'CosineAnnealingWarmRestarts':\n    scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)\nelif CFG.scheduler == 'ReduceLROnPlateau':\n    scheduler = ReduceLROnPlateauReduceLROnPlateau(optimizer, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)\nelif CFG.scheduler == 'CosineAnnealingLR':\n    scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f'Training Loop [{CFG.trn_fold}]...')\nfor fold, (tr_idx, val_idx) in enumerate(gkf.split(dir_df, groups=dir_df[dir_df.columns[0]].values)):\n    if fold != CFG.trn_fold: # Train only one fold\n        continue \n\n    trainloader, validloader = prepare_train_valid_dataloader(dir_df, [fold])\n\n    for epoch in range(CFG.epoch):\n        train_one_epoch(epoch, model, device, optimizer, scheduler, trainloader)\n        with torch.no_grad():\n            valid_one_epoch(epoch, model, device, optimizer, scheduler, validloader)\n        \n        #torch.save(model.state_dict(),f'FOLD-{fold}-EPOCH-{epoch}-model.pth')\n        \n    torch.save(model.state_dict(),f'FOLD-{fold}-model.pth')\n    \n    #del model, optimizer, scheduler, trainloader, validloader","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plot the results ","metadata":{}},{"cell_type":"code","source":"trainloader, validloader = prepare_train_valid_dataloader(dir_df, [4])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_result(validloader, n_sample=4):\n    img, mask = next(iter(validloader))\n    model.load_state_dict(torch.load(f'./FOLD-{CFG.trn_fold}-model.pth'))\n    pred_mask = model(img.to(device)).squeeze(1).cpu().detach()\n    \n    N = n_sample // 2\n    assert (N < CFG.batch_size)\n    plt.figure(figsize=(15, 20))\n    for i in range(n_sample):\n        plt.subplot(N, 4, 2*i+1)\n        plt.imshow(np.transpose(img[i], (1,2,0)))\n        plt.axis('off')\n        \n        plt.subplot(N, 4, 2*i+2)\n        plt.imshow(pred_mask[i])\n        plt.axis('off')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.no_grad():\n    plot_result(validloader)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## What you should do to improve the results:\n- Use data normalization\n- Use anthor implementation for TransUnet (Applied in V16)\n- Reduce the probability of the augmentations\n- Play with visionTrans hyperparameters","metadata":{}},{"cell_type":"markdown","source":"<CENTER><h3><span style='color:red'> UPVOTE </span> If you liked</h3>\n<h4>Code still under modifications, Stay tuned</h4></CENTER>","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}