{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.optim import Adam\nfrom torch.optim import lr_scheduler\nimport torch.nn.functional as F\nimport torchvision.transforms\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\nfrom torch.cuda.amp import autocast, GradScaler\n\nimport shutil\nimport pandas as pd\nfrom PIL import Image, ImageDraw\nimport os\nimport shutil\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold\nfrom sklearn import preprocessing\nfrom contextlib import contextmanager\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip,\n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout,\n    IAAAdditiveGaussianNoise, Transpose\n    )\nfrom albumentations.pytorch import ToTensorV2, ToTensor\nimport copy\nimport time\nfrom collections import defaultdict\nimport cv2\nimport matplotlib.pyplot as plt\nfrom glob import glob\nimport sys\nimport math","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass CFG:\n    num_workers = 4\n    size = 512\n    scheduler = 'CosineAnnealingLR'\n    segmentation_epochs = 15\n    segmentation_batch_size = 16\n    target_cols = ['ETT - Abnormal', 'ETT - Borderline', 'ETT - Normal',\n                   'NGT - Abnormal', 'NGT - Borderline', 'NGT - Incompletely Imaged', 'NGT - Normal',\n                   'CVC - Abnormal', 'CVC - Borderline', 'CVC - Normal',\n                   'Swan Ganz Catheter Present']\n    \nOUTPUT_DIR = './'\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)\n\nmodel_dir = \"weights\"\nbest_model_path = model_dir + \"/\" + \"best.pth\"\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ndef init_logger(log_file=OUTPUT_DIR+'train.log'):\n    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = init_logger()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyMiniUNet(nn.Module):\n    def __init__(self, n_channels, n_classes):\n        super(MyMiniUNet, self).__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n\n        self.inc = nn.Conv2d(n_channels, 32, kernel_size=3, padding=1)\n        self.down1 = MyDown(32, 64)\n        self.down2 = MyDown(64, 128)\n        self.down3 = MyDown(128, 256)\n\n        self.up1 = MyUp(384, 128)\n        self.up2 = MyUp(192, 64)\n        self.up3 = MyUp(96, 32)\n        self.outc = nn.Conv2d(32, n_classes, kernel_size=1)\n\n    def forward(self, x):\n\n        x1 = self.inc(x)\n        x2 = self.down1(x1)\n        x3 = self.down2(x2)\n        x4 = self.down3(x3)\n        x5 = self.up1(x4, x3)\n        x6 = self.up2(x5, x2)\n        x7 = self.up3(x6, x1)\n        output = self.outc(x7)\n        return output\n\n\nclass MyDown(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.my_conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.01)\n        )\n\n    def forward(self, x):\n        return self.my_conv(x)\n\n\nclass MyUp(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n\n        self.up = nn.Upsample(scale_factor=2, mode='nearest')\n\n        self.my_conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.01)\n        )\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n\n        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])\n        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])\n\n        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n                        diffY // 2, diffY - diffY // 2])\n\n        x1 = torch.cat([x2, x1], dim=1)\n        return self.my_conv(x1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainSegmentationDataset(Dataset):\n    def __init__(self, df, train_path, mask_path):\n        self.df = df\n        self.file_names = np.unique(df['StudyInstanceUID'].values)\n        self.transform = Compose([\n            Resize(CFG.size, CFG.size),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n        self.mask_transform = Compose([\n            Resize(CFG.size, CFG.size),\n            ToTensor(),\n        ])\n        self.train_path = train_path\n        self.mask_path = mask_path\n\n    def __len__(self):\n        return len(self.file_names)\n\n    def __getitem__(self, idx):\n        file_name = self.file_names[idx]\n        file_path = self.train_path + \"/\" + file_name + \".jpg\"\n        mfile_path = self.mask_path + \"/masks/\" + file_name + \".jpg\"\n        image = cv2.imread(file_path)\n        mask = cv2.imread(mfile_path, cv2.IMREAD_GRAYSCALE)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        label = self.mask_transform(image=mask)\n        image = self.transform(image=image)\n        image = image['image']\n        label = label['image']\n        label = label.reshape((1,) + label.shape)\n        return image, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef dice_loss(pred, target, smooth=1.):\n    pred = pred.contiguous()\n    target = target.contiguous()\n\n    intersection = (pred * target).sum(dim=2).sum(dim=2)\n\n    loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))\n\n    return loss.mean()\n\ndef calc_loss(pred, target, metrics, bce_weight=0.5):\n    bce = F.binary_cross_entropy_with_logits(pred, target)\n\n    pred = torch.sigmoid(pred)\n    dice = dice_loss(pred, target)\n\n    loss = bce * bce_weight + dice * (1 - bce_weight)\n\n    metrics['bce'] += bce.data.cpu().numpy() * target.size(0)\n    metrics['dice'] += dice.data.cpu().numpy() * target.size(0)\n    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)\n\n    return loss\n\ndef print_metrics(metrics, epoch_samples, phase):\n    outputs = []\n    for k in metrics.keys():\n        outputs.append(\"{}: {:4f}\".format(k, metrics[k] / epoch_samples))\n\n    LOGGER.info(\"{}: {}\".format(phase, \", \".join(outputs)))\n\n    \ndef train_model(model, optimizer, scheduler, dataloaders, num_epochs=25):\n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_loss = 1e10\n\n    for epoch in range(num_epochs):\n        LOGGER.info('Epoch {}/{}'.format(epoch, num_epochs - 1))\n        LOGGER.info('-' * 10)\n\n        since = time.time()\n\n        for phase in ['train', 'val']:\n            if phase == 'train':\n                scheduler.step()\n                for param_group in optimizer.param_groups:\n                    LOGGER.info(\"LR: \" + str(param_group['lr']))\n\n                model.train() \n            else:\n                model.eval()\n\n            metrics = defaultdict(float)\n            epoch_samples = 0\n\n            for inputs, labels in dataloaders[phase]:\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n\n                optimizer.zero_grad()\n\n                with torch.set_grad_enabled(phase == 'train'):\n                    outputs = model(inputs)\n                    loss = calc_loss(outputs, labels, metrics)\n\n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n\n                epoch_samples += inputs.size(0)\n\n            print_metrics(metrics, epoch_samples, phase)\n            epoch_loss = metrics['loss'] / epoch_samples\n\n            if phase == 'val' and epoch_loss < best_loss:\n                LOGGER.info(\"saving best model\")\n                best_loss = epoch_loss\n                best_model_wts = copy.deepcopy(model.state_dict())\n                torch.save(best_model_wts, best_model_path)\n\n        time_elapsed = time.time() - since\n        LOGGER.info('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n\n    LOGGER.info('Best val loss: {:4f}'.format(best_loss))\n\n    model.load_state_dict(best_model_wts)\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists(model_dir):\n    os.makedirs(model_dir)\ntrain_annotations = pd.read_csv(\"../input/ranzcr-clip-catheter-line-classification/train_annotations.csv\")\nimages_path = \"../input/ranzcr-clip-catheter-line-classification/train\"\nmask_path = \"../input/generate-masks-ranzr\"\ntraining, validation = train_test_split(train_annotations, test_size=0.2)\ntraining = training.reset_index()\nvalidation = validation.reset_index()\ntrain_dataset = TrainSegmentationDataset(training, images_path, mask_path)\nvalid_dataset = TrainSegmentationDataset(validation, images_path, mask_path)\ntrain_loader = DataLoader(train_dataset,batch_size=CFG.segmentation_batch_size,shuffle=True,num_workers=CFG.num_workers, pin_memory=True, drop_last=True)\nvalid_loader = DataLoader(valid_dataset,batch_size=CFG.segmentation_batch_size,shuffle=False,num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\nmodel = MyMiniUNet(3, 1)\noptimizer_ft = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)\ndataloaders = {\"train\": train_loader, \"val\": valid_loader}\nexp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=30, gamma=0.1)\nmodel = model.to(device)\nLOGGER.info(\"****** starting training of segmentation model ******\")\ntrained_model = train_model(model, optimizer_ft, exp_lr_scheduler,dataloaders, CFG.segmentation_epochs)\nLOGGER.info(\"****** ended training of segmentation model ******\")","execution_count":null,"outputs":[]}],"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}