{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport pdb\nimport time\nimport warnings\nimport random\nimport zipfile\nimport pydicom\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm_notebook as tqdm\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom sklearn.model_selection import train_test_split\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nimport torch.optim as optim\nimport torch.backends.cudnn as cudnn\nfrom torch.utils.data import DataLoader, Dataset, sampler\nfrom matplotlib import pyplot as plt\nimport seaborn as sns\nfrom albumentations import (HorizontalFlip, VerticalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise)\nfrom albumentations.pytorch import ToTensor\nwarnings.simplefilter(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/qubvel/segmentation_models.pytorch > /dev/null 2>&1 # Install segmentations_models.pytorch, with no bash output\nimport segmentation_models_pytorch as smp","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Preprocessing & Utility functional","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def set_seed(seed=2**3):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.backends.cudnn.deterministic = True\nset_seed(121)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transforms(phase, size, mean, std):\n    list_transforms = []\n    if phase == \"train\":\n        list_transforms.extend(\n            [\n                HorizontalFlip(p=0.5),\n                VerticalFlip(p=0.5),\n                ShiftScaleRotate(\n                    shift_limit=0,  # no resizing\n                    scale_limit=0.1,\n                    rotate_limit=5, # rotate\n                    p=0.5,\n                    border_mode=cv2.BORDER_CONSTANT\n                ),\n            ]\n        )\n    list_transforms.extend(\n        [\n            Resize(size, size),\n            Normalize(mean=mean, std=std, p=1),\n            ToTensor(),\n        ]\n    )\n\n    list_trfms = Compose(list_transforms)\n    return list_trfms\n\ndef train_transform_creator():\n    return get_transforms(\n        'train',\n        512,\n        0.,\n        1.\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class OSICDataset(Dataset):\n    def __init__(self, fnames, img_folder, mask_folder, transforms_creator):\n        self.img_folder = img_folder\n        self.mask_folder = mask_folder\n        self.transforms = transforms_creator()\n        self.fnames = fnames\n\n    def __getitem__(self, idx):\n        image_id = self.fnames[idx]\n        d = pydicom.dcmread(os.path.join(self.img_folder, image_id))\n        image = (d.pixel_array - d.RescaleIntercept) / (d.RescaleSlope * 1000)\n        image[image < 1.5] = 0.\n        mask = cv2.imread(os.path.join(self.mask_folder, f'{image_id[:-4]}.jpg'), cv2.IMREAD_GRAYSCALE)\n        image = np.dstack([image, image, image])\n        augmented = self.transforms(image=image, mask=mask)\n        image = augmented['image']\n        mask = augmented['mask']\n        return image, mask\n\n    def __len__(self):\n        return len(self.fnames)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's make stratify train/validation set by mask area.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"label = []\nfnames = []\ntrain = pd.read_csv('../input/osic-pulmonary-fibrosis-progression/train.csv')\n\nfor i, p in tqdm(enumerate(train.Patient.unique())):\n    for k in os.listdir(f'../input/osic-pulmonary-fibrosis-progression/train/{p}/'):\n        m = cv2.imread(f'../input/osic-pulmonary-fibrosis-progression-lungs-mask/mask_clear/mask_clear/{p}/{k[:-4]}.jpg')\n        if m is None:\n            continue\n        fnames.append(f'{p}/{k}')\n        label.append(m.sum())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_label = np.array(label) // 1000000\n_label[_label > 40] = 40\nsns.distplot(_label)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tr_fnames, vl_fnames, _, _ = train_test_split(fnames, _label, train_size=0.85, random_state=41, shuffle=True, stratify=_label)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tr_image_dataset = OSICDataset(tr_fnames, \n                            '../input/osic-pulmonary-fibrosis-progression/train', \n                            '../input/osic-pulmonary-fibrosis-progression-lungs-mask/mask_clear/mask_clear', \n                            train_transform_creator)\nvl_image_dataset = OSICDataset(vl_fnames, \n                            '../input/osic-pulmonary-fibrosis-progression/train', \n                            '../input/osic-pulmonary-fibrosis-progression-lungs-mask/mask_clear/mask_clear', \n                            train_transform_creator)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tr_dataloader = DataLoader(\n        tr_image_dataset,\n        batch_size=16,\n        num_workers=5,\n        pin_memory=True,\n        shuffle=True,\n    )\n\nvl_dataloader = DataLoader(\n        vl_image_dataset,\n        batch_size=16,\n        num_workers=5,\n        pin_memory=True,\n        shuffle=True,\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch = next(iter(tr_dataloader)) # get a batch from the dataloader\nimages, masks = batch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = random.choice(range(32))\nplt.figure(figsize=(20,20))\nplt.imshow(images[idx].reshape((512, 512, 3)), cmap='bone')\nplt.imshow(masks[idx][0], alpha=0.3, cmap='Reds')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Losses & Metrics","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def dice_loss(input, target):\n    input = torch.sigmoid(input)\n    smooth = 1.0\n    iflat = input.view(-1)\n    tflat = target.view(-1)\n    intersection = (iflat * tflat).sum()\n    return ((2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))\n\n\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma):\n        super().__init__()\n        self.gamma = gamma\n\n    def forward(self, input, target):\n        if not (target.size() == input.size()):\n            raise ValueError(\"Target size ({}) must be the same as input size ({})\"\n                             .format(target.size(), input.size()))\n        max_val = (-input).clamp(min=0)\n        loss = input - input * target + max_val + \\\n            ((-max_val).exp() + (-input - max_val).exp()).log()\n        invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))\n        loss = (invprobs * self.gamma).exp() * loss\n        return loss.mean()\n\n\nclass MixedLoss(nn.Module):\n    def __init__(self, alpha, gamma):\n        super().__init__()\n        self.alpha = alpha\n        self.focal = FocalLoss(gamma)\n\n    def forward(self, input, target):\n        loss = self.alpha*self.focal(input, target) - torch.log(dice_loss(input, target))\n        return loss.mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def predict(X, threshold):\n    X_p = np.copy(X)\n    preds = (X_p > threshold).astype('uint8')\n    return preds\n\ndef metric(probability, truth, threshold=0.5, reduction='none'):\n    '''Calculates dice of positive and negative images seperately'''\n    '''probability and truth must be torch tensors'''\n    batch_size = len(truth)\n    with torch.no_grad():\n        probability = probability.view(batch_size, -1)\n        truth = truth.view(batch_size, -1)\n        assert(probability.shape == truth.shape)\n\n        p = (probability > threshold).float()\n        t = (truth > 0.5).float()\n\n        t_sum = t.sum(-1)\n        p_sum = p.sum(-1)\n        neg_index = torch.nonzero(t_sum == 0)\n        pos_index = torch.nonzero(t_sum >= 1)\n\n        dice_neg = (p_sum == 0).float()\n        dice_pos = 2 * (p*t).sum(-1)/((p+t).sum(-1))\n\n        dice_neg = dice_neg[neg_index]\n        dice_pos = dice_pos[pos_index]\n        dice = torch.cat([dice_pos, dice_neg])\n\n        num_neg = len(neg_index)\n        num_pos = len(pos_index)\n\n    return dice, dice_neg, dice_pos, num_neg, num_pos\n\nclass Meter:\n    '''A meter to keep track of iou and dice scores throughout an epoch'''\n    def __init__(self, phase, epoch):\n        self.base_threshold = 0.5\n        self.base_dice_scores = []\n        self.dice_neg_scores = []\n        self.dice_pos_scores = []\n        self.iou_scores = []\n\n    def update(self, targets, outputs):\n        probs = torch.sigmoid(outputs)\n        dice, dice_neg, dice_pos, _, _ = metric(probs, targets, self.base_threshold)\n        self.base_dice_scores.extend(dice)\n        self.dice_pos_scores.extend(dice_pos)\n        self.dice_neg_scores.extend(dice_neg)\n        preds = predict(probs, self.base_threshold)\n        iou = compute_iou_batch(preds, targets, classes=[1])\n        self.iou_scores.append(iou)\n\n    def get_metrics(self):\n        dice = np.nanmean(self.base_dice_scores)\n        dice_neg = np.nanmean(self.dice_neg_scores)\n        dice_pos = np.nanmean(self.dice_pos_scores)\n        dices = [dice, dice_neg, dice_pos]\n        iou = np.nanmean(self.iou_scores)\n        return dices, iou\n\ndef epoch_log(phase, epoch, epoch_loss, meter, start):\n    '''logging the metrics at the end of an epoch'''\n    dices, iou = meter.get_metrics()\n    dice, dice_neg, dice_pos = dices\n    print(\"Loss: %0.4f | dice: %0.4f | dice_neg: %0.4f | dice_pos: %0.4f | IoU: %0.4f\" % (epoch_loss, dice, dice_neg, dice_pos, iou))\n    return dice, iou\n\ndef compute_ious(pred, label, classes, ignore_index=255, only_present=True):\n    '''computes iou for one ground truth mask and predicted mask'''\n    pred[label == ignore_index] = 0\n    ious = []\n    for c in classes:\n        label_c = label == c\n        if only_present and np.sum(label_c) == 0:\n            ious.append(np.nan)\n            continue\n        pred_c = pred == c\n        intersection = np.logical_and(pred_c, label_c).sum()\n        union = np.logical_or(pred_c, label_c).sum()\n        if union != 0:\n            ious.append(intersection / union)\n    return ious if ious else [1]\n\n\ndef compute_iou_batch(outputs, labels, classes=None):\n    '''computes mean iou for a batch of ground truth masks and predicted masks'''\n    ious = []\n    preds = np.copy(outputs) # copy is imp\n    labels = np.array(labels) # tensor to np\n    for pred, label in zip(preds, labels):\n        ious.append(np.nanmean(compute_ious(pred, label, classes)))\n    iou = np.nanmean(ious)\n    return iou","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[](https://miro.medium.com/max/1000/1*6WrSlLJ69fRNB2RER5M_Vw.png)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model = smp.Unet(\"resnet34\", encoder_weights=\"imagenet\", activation=None)\nmodel","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training UNet","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Trainer(object):\n    '''This class takes care of training and validation of our model'''\n    def __init__(self, model):\n        self.fold = 1\n        self.total_folds = 5\n        self.num_workers = 6\n        self.batch_size = {\"train\": 16, \"val\": 16}\n        self.accumulation_steps = 32 // self.batch_size['train']\n        self.lr = 1e-3\n        self.num_epochs = 7\n        self.best_loss = float(\"inf\")\n        self.phases = [\"train\", \"val\"]\n        self.device = torch.device(\"cuda:0\")\n        torch.set_default_tensor_type(\"torch.cuda.FloatTensor\")\n        self.net = model\n        self.criterion = FocalLoss(1.5)\n        self.optimizer = optim.Adam(self.net.parameters(), lr=self.lr)\n        self.scheduler = ReduceLROnPlateau(self.optimizer, mode=\"min\", patience=3, verbose=True)\n        self.net = self.net.to(self.device)\n        cudnn.benchmark = True\n        self.dataloaders = {\n            phase: loader\n            for phase, loader in zip(self.phases, [tr_dataloader, vl_dataloader])\n        }\n        self.losses = {phase: [] for phase in self.phases}\n        self.iou_scores = {phase: [] for phase in self.phases}\n        self.dice_scores = {phase: [] for phase in self.phases}\n        \n    def forward(self, images, targets):\n        images = images.to(self.device)\n        masks = targets.to(self.device)\n        outputs = self.net(images)\n        loss = self.criterion(outputs, masks)\n        return loss, outputs\n\n    def iterate(self, epoch, phase):\n        meter = Meter(phase, epoch)\n        start = time.strftime(\"%H:%M:%S\")\n        print(f\"Starting epoch: {epoch} | phase: {phase} | ⏰: {start}\")\n        batch_size = self.batch_size[phase]\n        self.net.train(phase == \"train\")\n        dataloader = self.dataloaders[phase]\n        running_loss = 0.0\n        total_batches = len(dataloader)\n\n        self.optimizer.zero_grad()\n        for itr, batch in enumerate(dataloader):\n            images, targets = batch\n            loss, outputs = self.forward(images, targets)\n            loss = loss / self.accumulation_steps\n            if phase == \"train\":\n                loss.backward()\n                if (itr + 1 ) % self.accumulation_steps == 0:\n                    self.optimizer.step()\n                    self.optimizer.zero_grad()\n            running_loss += loss.item()\n            outputs = outputs.detach().cpu()\n            meter.update(targets, outputs)\n\n        epoch_loss = (running_loss * self.accumulation_steps) / total_batches\n        dice, iou = epoch_log(phase, epoch, epoch_loss, meter, start)\n        self.losses[phase].append(epoch_loss)\n        self.dice_scores[phase].append(dice)\n        self.iou_scores[phase].append(iou)\n        torch.cuda.empty_cache()\n        return epoch_loss\n\n    def start(self):\n        for epoch in range(self.num_epochs):\n            self.iterate(epoch, \"train\")\n            state = {\n                \"epoch\": epoch,\n                \"best_loss\": self.best_loss,\n                \"state_dict\": self.net.state_dict(),\n                \"optimizer\": self.optimizer.state_dict(),\n            }\n            val_loss = self.iterate(epoch, \"val\")\n            self.scheduler.step(val_loss)\n            if val_loss < self.best_loss:\n                print(\"******** New optimal found, saving state ********\")\n                state[\"best_loss\"] = self.best_loss = val_loss\n                torch.save(self.net, 'model.pth')\n            print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model_trainer = Trainer(model)\nmodel_trainer.start()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\ndel tr_dataloader, vl_dataloader, Trainer, tr_image_dataset, vl_image_dataset\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torch.load(\"model.pth\")\nmodel","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prediction","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def zip_and_remove(path):\n    ziph = zipfile.ZipFile(f'{path}.zip', 'w', zipfile.ZIP_DEFLATED)\n    \n    for root, dirs, files in os.walk(path):\n        for file in tqdm(files):\n            file_path = os.path.join(root, file)\n            ziph.write(file_path)\n            os.remove(file_path)\n    \n    ziph.close()\n    \nif not os.path.exists('train_mask_unet_prob/'):\n    os.mkdir('train_mask_unet_prob/')\n    \nif not os.path.exists('train_mask_unet/'):\n    os.mkdir('train_mask_unet/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/osic-pulmonary-fibrosis-progression/train.csv')\ntrain_data = {}\nfor p in train.Patient.values:\n    train_data[p] = os.listdir(f'../input/osic-pulmonary-fibrosis-progression/train/{p}/')\n    \nkeys = [k for k in list(train_data.keys()) if k not in ['ID00011637202177653955184', 'ID00052637202186188008618']]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms = get_transforms('valid', 512, 0., 1.)\n\nfor k in tqdm(keys, total=len(keys)):\n    x = []\n    if not os.path.exists('train_mask_unet_prob/' + k):\n        os.mkdir('train_mask_unet_prob/' + k)\n        \n    if not os.path.exists('train_mask_unet/' + k):\n        os.mkdir('train_mask_unet/' + k)\n        \n    for i in train_data[k]:\n        d =  pydicom.dcmread(f'../input/osic-pulmonary-fibrosis-progression/train/{k}/{i}')\n        image = (d.pixel_array - d.RescaleIntercept) / (d.RescaleSlope * 1000)\n        image[image < 1.5] = 0.\n        image = np.dstack([image, image, image])\n        image = transforms(image=image)['image'].reshape((1, 3, 512, 512)).to('cuda')\n        \n        mask = model(image).detach().to('cpu').numpy()[0, 0, ...]\n        \n        cv2. imwrite('train_mask_unet_prob/' + k + f'/{i[:-4]}' + '.jpg', mask)\n        cv2. imwrite('train_mask_unet/' + k + f'/{i[:-4]}' + '.jpg', np.uint8(mask > 0.5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('test_mask_unet_prob/'):\n    os.mkdir('test_mask_unet_prob/')\n    \nif not os.path.exists('test_mask_unet/'):\n    os.mkdir('test_mask_unet/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = pd.read_csv('../input/osic-pulmonary-fibrosis-progression/test.csv')\ntest_data = {}\nfor p in test.Patient.values:\n    test_data[p] = os.listdir(f'../input/osic-pulmonary-fibrosis-progression/test/{p}/')\n    \nkeys = [k for k in list(test_data.keys()) if k not in ['ID00011637202177653955184', 'ID00052637202186188008618']]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms = get_transforms('valid', 512, 0., 1.)\n\nfor k in tqdm(keys, total=len(keys)):\n    x = []\n    if not os.path.exists('test_mask_unet_prob/' + k):\n        os.mkdir('test_mask_unet_prob/' + k)\n        \n    if not os.path.exists('test_mask_unet/' + k):\n        os.mkdir('test_mask_unet/' + k)\n        \n    for i in train_data[k]:\n        d =  pydicom.dcmread(f'../input/osic-pulmonary-fibrosis-progression/test/{k}/{i}')\n        image = (d.pixel_array - d.RescaleIntercept) / (d.RescaleSlope * 1000)\n        image[image < 1.5] = 0.\n        image = np.dstack([image, image, image])\n        image = transforms(image=image)['image'].reshape((1, 3, 512, 512)).to('cuda')\n        \n        mask = model(image).detach().to('cpu').numpy()[0, 0, ...]\n        \n        cv2. imwrite('test_mask_unet_prob/' + k + f'/{i[:-4]}' + '.jpg', mask)\n        cv2. imwrite('test_mask_unet/' + k + f'/{i[:-4]}' + '.jpg', np.uint8(mask > 0.5))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Result Visualization","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"masks = os.listdir('train_mask_unet_prob/ID00408637202308839708961')\nmasks = sorted(masks, key=lambda x: int(x[:-4]))\nlen(masks)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_, axs = plt.subplots(6, 6, figsize=(24, 24))\naxs = axs.flatten()\nfor m, ax in zip(masks, axs):\n    img = cv2.imread(os.path.join('train_mask_unet_prob/ID00408637202308839708961', m), cv2.IMREAD_GRAYSCALE)\n    ax.imshow(img)\n    ax.axis('off')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"zip_and_remove('train_mask_unet_prob')\nzip_and_remove('train_mask_unet')\nzip_and_remove('test_mask_unet_prob')\nzip_and_remove('test_mask_unet')","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}