{"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":"code","source":"!unzip ../input/data-science-bowl-2018/stage1_train.zip > null","metadata":{"execution":{"iopub.status.busy":"2022-01-03T13:34:51.544895Z","iopub.execute_input":"2022-01-03T13:34:51.545763Z","iopub.status.idle":"2022-01-03T13:34:55.259399Z","shell.execute_reply.started":"2022-01-03T13:34:51.545641Z","shell.execute_reply":"2022-01-03T13:34:55.258502Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from albumentations import pytorch\nimport torch\nfrom torch.utils.data import DataLoader, Dataset, random_split\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torchvision.transforms import CenterCrop\n\nimport numpy as np\nfrom cv2 import cv2, transform\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport pytorch_lightning as pl\nfrom pathlib import Path","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-01-03T13:34:55.261660Z","iopub.execute_input":"2022-01-03T13:34:55.262215Z","iopub.status.idle":"2022-01-03T13:35:02.987921Z","shell.execute_reply.started":"2022-01-03T13:34:55.262176Z","shell.execute_reply":"2022-01-03T13:35:02.986916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 32","metadata":{"execution":{"iopub.status.busy":"2022-01-03T13:35:02.989736Z","iopub.execute_input":"2022-01-03T13:35:02.990038Z","iopub.status.idle":"2022-01-03T13:35:02.995723Z","shell.execute_reply.started":"2022-01-03T13:35:02.989984Z","shell.execute_reply":"2022-01-03T13:35:02.993277Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, c_in, c_out):\n        super(DoubleConv, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=3),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(in_channels=c_out, out_channels=c_out, kernel_size=3),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass DoubleConvSame(nn.Module):\n    def __init__(self, c_in, c_out):\n        super(DoubleConvSame, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(in_channels=c_out, out_channels=c_out, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass UNet_2(nn.Module):\n    def __init__(self, c_in, c_out):\n        super(UNet_2, self).__init__()\n\n        self.conv1 = DoubleConvSame(c_in=c_in, c_out=64)\n        self.conv2 = DoubleConvSame(c_in=64, c_out=128)\n        self.conv3 = DoubleConvSame(c_in=128, c_out=256)\n        self.conv4 = DoubleConvSame(c_in=256, c_out=512)\n        self.conv5 = DoubleConvSame(c_in=512, c_out=1024)\n\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        self.up1 = nn.ConvTranspose2d(\n            in_channels=1024, out_channels=512, kernel_size=2, stride=2\n        )\n        self.up2 = nn.ConvTranspose2d(\n            in_channels=512, out_channels=256, kernel_size=2, stride=2\n        )\n        self.up3 = nn.ConvTranspose2d(\n            in_channels=256, out_channels=128, kernel_size=2, stride=2\n        )\n        self.up4 = nn.ConvTranspose2d(\n            in_channels=128, out_channels=64, kernel_size=2, stride=2\n        )\n\n        self.up_conv1 = DoubleConvSame(c_in=1024, c_out=512)\n        self.up_conv2 = DoubleConvSame(c_in=512, c_out=256)\n        self.up_conv3 = DoubleConvSame(c_in=256, c_out=128)\n        self.up_conv4 = DoubleConvSame(c_in=128, c_out=64)\n\n        self.conv_1x1 = nn.Conv2d(in_channels=64, out_channels=c_out, kernel_size=1)\n\n    def forward(self, x):\n        \"\"\"ENCODER\"\"\"\n\n        c1 = self.conv1(x)\n        p1 = self.pool(c1)\n\n        c2 = self.conv2(p1)\n        p2 = self.pool(c2)\n\n        c3 = self.conv3(p2)\n        p3 = self.pool(c3)\n\n        c4 = self.conv4(p3)\n        p4 = self.pool(c4)\n        \"\"\"BOTTLE-NECK\"\"\"\n\n        c5 = self.conv5(p4)\n        \"\"\"DECODER\"\"\"\n\n        u1 = self.up1(c5)\n        cat1 = torch.cat([u1, c4], dim=1)\n        uc1 = self.up_conv1(cat1)\n\n        u2 = self.up2(uc1)\n        cat2 = torch.cat([u2, c3], dim=1)\n        uc2 = self.up_conv2(cat2)\n\n        u3 = self.up3(uc2)\n        cat3 = torch.cat([u3, c2], dim=1)\n        uc3 = self.up_conv3(cat3)\n\n        u4 = self.up4(uc3)\n        cat4 = torch.cat([u4, c1], dim=1)\n        uc4 = self.up_conv4(cat4)\n\n        outputs = self.conv_1x1(uc4)\n\n        return outputs\n\n\nclass UNet_OG(nn.Module):\n    def __init__(self, c_in, c_out):\n        super(UNet_OG, self).__init__()\n\n        self.conv1 = DoubleConv(c_in=c_in, c_out=64)\n        self.conv2 = DoubleConv(c_in=64, c_out=128)\n        self.conv3 = DoubleConv(c_in=128, c_out=256)\n        self.conv4 = DoubleConv(c_in=256, c_out=512)\n        self.conv5 = DoubleConv(c_in=512, c_out=1024)\n\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        self.up1 = nn.ConvTranspose2d(\n            in_channels=1024, out_channels=512, kernel_size=2, stride=2\n        )\n        self.up2 = nn.ConvTranspose2d(\n            in_channels=512, out_channels=256, kernel_size=2, stride=2\n        )\n        self.up3 = nn.ConvTranspose2d(\n            in_channels=256, out_channels=128, kernel_size=2, stride=2\n        )\n        self.up4 = nn.ConvTranspose2d(\n            in_channels=128, out_channels=64, kernel_size=2, stride=2\n        )\n\n        self.up_conv1 = DoubleConv(c_in=1024, c_out=512)\n        self.up_conv2 = DoubleConv(c_in=512, c_out=256)\n        self.up_conv3 = DoubleConv(c_in=256, c_out=128)\n        self.up_conv4 = DoubleConv(c_in=128, c_out=64)\n\n        self.conv_1x1 = nn.Conv2d(in_channels=64, out_channels=c_out, kernel_size=1)\n\n    def crop_tensor(self, up_tensor, target_tensor):\n        _, _, H, W = up_tensor.shape\n\n        x = CenterCrop(size=(H, W))(target_tensor)\n\n        return x\n\n    def forward(self, x):\n        \"\"\"ENCODER\"\"\"\n\n        c1 = self.conv1(x)\n        p1 = self.pool(c1)\n\n        c2 = self.conv2(p1)\n        p2 = self.pool(c2)\n\n        c3 = self.conv3(p2)\n        p3 = self.pool(c3)\n\n        c4 = self.conv4(p3)\n        p4 = self.pool(c4)\n        \"\"\"BOTTLE-NECK\"\"\"\n\n        c5 = self.conv5(p4)\n        \"\"\"DECODER\"\"\"\n\n        u1 = self.up1(c5)\n        crop1 = self.crop_tensor(u1, c4)\n        cat1 = torch.cat([u1, crop1], dim=1)\n        uc1 = self.up_conv1(cat1)\n\n        u2 = self.up2(uc1)\n        crop2 = self.crop_tensor(u2, c3)\n        cat2 = torch.cat([u2, crop2], dim=1)\n        uc2 = self.up_conv2(cat2)\n\n        u3 = self.up3(uc2)\n        crop3 = self.crop_tensor(u3, c2)\n        cat3 = torch.cat([u3, crop3], dim=1)\n        uc3 = self.up_conv3(cat3)\n\n        u4 = self.up4(uc3)\n        crop4 = self.crop_tensor(u4, c1)\n        cat4 = torch.cat([u4, crop4], dim=1)\n        uc4 = self.up_conv4(cat4)\n\n        outputs = self.conv_1x1(uc4)\n\n        return outputs\n\nclass AttentionBlock(nn.Module):\n    \"\"\"\n    Class for creating Attention module\n    Takes in gating signal `g` and `x`\n    \"\"\"\n\n    def __init__(self, g_chl, x_chl):\n        super(AttentionBlock, self).__init__()\n\n        inter_shape = x_chl // 4\n\n        # Conv 1x1 with stride 2 for `x`\n        self.conv_x = nn.Conv2d(\n            in_channels=x_chl,\n            out_channels=inter_shape,\n            kernel_size=1,\n            stride=2,\n        )\n\n        # Conv 1x1 with stride 1 for `g` (gating signal)\n        self.conv_g = nn.Conv2d(\n            in_channels=g_chl,\n            out_channels=inter_shape,\n            kernel_size=1,\n            stride=1,\n        )\n\n        # Conv 1x1 for `psi` the output after `g` + `x`\n        self.psi = nn.Conv2d(\n            in_channels=2 * inter_shape,\n            out_channels=1,\n            kernel_size=1,\n            stride=1,\n        )\n\n        # For upsampling the attention output to size of `x`\n        self.upsample = nn.Upsample(scale_factor=2)\n\n    def forward(self, g, x):\n\n        # perform the convs on `x` and `g`\n        theta_x = self.conv_x(x)\n        gate = self.conv_g(g)\n\n        # `theta_x` + `gate`\n        add = torch.cat([self.conv_x(x), self.conv_g(g)], axis=1)\n\n        # ReLU on the add operation\n        relu = torch.relu(add)\n\n        # the 1x1 Conv\n        psi = self.psi(relu)\n\n        # Sigmoid the squash the outputs/attention weights\n        sig = torch.sigmoid(psi)\n\n        # Upsample to original size of `x` to perform multiplication\n        upsample = self.upsample(sig)\n\n        # return the attention weights!\n        return upsample\n\n\nclass AttentionUNet(nn.Module):\n    def __init__(self, c_in, c_out):\n        super(AttentionUNet, self).__init__()\n\n        self.conv1 = DoubleConvSame(c_in=c_in, c_out=64)\n        self.conv2 = DoubleConvSame(c_in=64, c_out=128)\n        self.conv3 = DoubleConvSame(c_in=128, c_out=256)\n        self.conv4 = DoubleConvSame(c_in=256, c_out=512)\n        self.conv5 = DoubleConvSame(c_in=512, c_out=1024)\n\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        self.attn1 = AttentionBlock(1024, 512)\n        self.attn2 = AttentionBlock(512, 256)\n        self.attn3 = AttentionBlock(256, 128)\n        self.attn4 = AttentionBlock(128, 64)\n\n        self.up1 = nn.ConvTranspose2d(\n            in_channels=1024, out_channels=512, kernel_size=2, stride=2\n        )\n        self.up2 = nn.ConvTranspose2d(\n            in_channels=512, out_channels=256, kernel_size=2, stride=2\n        )\n        self.up3 = nn.ConvTranspose2d(\n            in_channels=256, out_channels=128, kernel_size=2, stride=2\n        )\n        self.up4 = nn.ConvTranspose2d(\n            in_channels=128, out_channels=64, kernel_size=2, stride=2\n        )\n\n        self.up_conv1 = DoubleConvSame(c_in=1024, c_out=512)\n        self.up_conv2 = DoubleConvSame(c_in=512, c_out=256)\n        self.up_conv3 = DoubleConvSame(c_in=256, c_out=128)\n        self.up_conv4 = DoubleConvSame(c_in=128, c_out=64)\n\n        self.conv_1x1 = nn.Conv2d(in_channels=64, out_channels=c_out, kernel_size=1)\n\n    def forward(self, x):\n        \"\"\"ENCODER\"\"\"\n\n        c1 = self.conv1(x)\n        p1 = self.pool(c1)\n\n        c2 = self.conv2(p1)\n        p2 = self.pool(c2)\n\n        c3 = self.conv3(p2)\n        p3 = self.pool(c3)\n\n        c4 = self.conv4(p3)\n        p4 = self.pool(c4)\n\n        \"\"\"BOTTLE-NECK\"\"\"\n\n        c5 = self.conv5(p4)\n\n        \"\"\"DECODER - WITH ATTENTION\"\"\"\n\n        att1 = self.attn1(c5, c4)\n        u1 = self.up1(c5)\n        mult1 = torch.multiply(att1, u1)\n        cat1 = torch.cat([mult1, c4], dim=1)\n        uc1 = self.up_conv1(cat1)\n\n        att2 = self.attn2(uc1, c3)\n        u2 = self.up2(uc1)\n        mult2 = torch.multiply(att2, u2)\n        cat2 = torch.cat([mult2, c3], dim=1)\n        uc2 = self.up_conv2(cat2)\n\n        att3 = self.attn3(uc2, c2)\n        u3 = self.up3(uc2)\n        mult3 = torch.multiply(att3, u3)\n        cat3 = torch.cat([mult3, c2], dim=1)\n        uc3 = self.up_conv3(cat3)\n\n        att4 = self.attn4(uc3, c1)\n        u4 = self.up4(uc3)\n        mult4 = torch.multiply(att4, u4)\n        cat4 = torch.cat([mult4, c1], dim=1)\n        uc4 = self.up_conv4(cat4)\n\n        outputs = self.conv_1x1(uc4)\n\n        return outputs","metadata":{"execution":{"iopub.status.busy":"2022-01-03T13:35:02.997868Z","iopub.execute_input":"2022-01-03T13:35:02.998257Z","iopub.status.idle":"2022-01-03T13:35:03.055685Z","shell.execute_reply.started":"2022-01-03T13:35:02.998220Z","shell.execute_reply":"2022-01-03T13:35:03.054808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class NucleiData(Dataset):\n    def __init__(\n        self, data_dir=\"/kaggle/working/\", transforms=None\n    ):\n        train_dir = Path(data_dir)\n        self.images = list(train_dir.glob(\"*/images/*.png\"))\n        self.masks = list(train_dir.glob(\"*/masks/*.*\"))\n        self.transforms = transforms\n\n    def __getitem__(self, idx):\n        image = cv2.imread(self.images[idx].as_posix(), cv2.IMREAD_COLOR)\n        mask = self.get_mask(self.images[idx].parent.parent.glob(\"masks/*.*\"))\n\n        if self.transforms is not None:\n            transform = self.transforms(image=image)\n        \n        transformed_image = transform[\"image\"]\n        transformed_mask = ToTensorV2()(image=mask)\n\n        return transformed_image, transformed_mask['image']\n\n    def get_mask(self, masks_gen):\n        H, W = 256, 256\n        target_mask = np.zeros((H, W, 1), dtype=np.uint8)\n        for mask in masks_gen:\n            curr_mask = cv2.imread(mask.as_posix(), cv2.IMREAD_GRAYSCALE)\n            transform = A.Resize(height=H, width=W)(image=curr_mask)\n            mask_ = np.expand_dims(transform[\"image\"], axis=-1)\n            target_mask = np.maximum(target_mask, mask_)\n        return target_mask\n\n    def __len__(self):\n        return len(self.images)\n\n\nclass NucleiDataModule(pl.LightningDataModule):\n    def __init__(self):\n        super().__init__()\n        self.image_transforms = A.Compose(\n            [A.Resize(256, 256), A.Normalize(), A.pytorch.ToTensorV2()]\n        )\n        self.dims = (3, 256, 256)\n\n    def setup(self, stage) -> None:\n        if stage == \"fit\" or stage is None:\n            data = NucleiData(transforms=self.image_transforms)\n            lengths = [int(len(data) * 0.8), int(len(data) * 0.2)]\n            self.train_data, self.val_data = random_split(data, lengths)\n\n    def train_dataloader(self):\n        return DataLoader(self.train_data, batch_size=BATCH_SIZE)\n\n    def val_dataloader(self):\n        return DataLoader(self.val_data, batch_size=BATCH_SIZE)\n\nclass LitNuclei(pl.LightningModule):\n    def __init__(self):\n        super(LitNuclei, self).__init__()\n        self.model = AttentionUNet(3, 1)\n        self.loss = nn.BCEWithLogitsLoss()\n        \n    def configure_optimizers(self):\n        return optim.Adam(self.model.parameters())\n    \n    def forward(self, x):\n        return self.model(x)\n    \n    def training_step(self, batch, batch_idx):\n        image, mask = batch\n        \n        preds = self.forward(image)\n        \n        loss = F.binary_cross_entropy_with_logits(\n            input=preds, target=mask.float()\n        )\n        \n        self.log('train_loss', loss)\n        \n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        image, mask = batch\n        \n        preds = self.forward(image)\n        \n        loss = F.binary_cross_entropy_with_logits(\n            input=preds, target=mask.float()\n        )\n        \n        self.log('val_loss', loss)\n        \n        return loss","metadata":{"execution":{"iopub.status.busy":"2022-01-03T13:35:03.058509Z","iopub.execute_input":"2022-01-03T13:35:03.058702Z","iopub.status.idle":"2022-01-03T13:35:03.078769Z","shell.execute_reply.started":"2022-01-03T13:35:03.058681Z","shell.execute_reply":"2022-01-03T13:35:03.077195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = LitNuclei()\ndm = NucleiDataModule()\n\ntrainer = pl.Trainer(\n    checkpoint_callback=True,\n    logger=True,\n    max_epochs=5, gpus=1, \n)\n\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2022-01-03T13:35:03.079884Z","iopub.execute_input":"2022-01-03T13:35:03.080162Z","iopub.status.idle":"2022-01-03T13:39:07.273403Z","shell.execute_reply.started":"2022-01-03T13:35:03.080128Z","shell.execute_reply":"2022-01-03T13:39:07.272597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}