{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install torchsummary","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport cv2\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport torch\nimport torchvision\nfrom torchsummary import summary\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms, models\nfrom torch.optim import lr_scheduler, Adam\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, \n    HueSaturationValue, CoarseDropout, ToGray\n    )\nfrom albumentations.pytorch import ToTensorV2\n\ndevice = (\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndevice","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_folder = \"/kaggle/input/hpa-single-cell-image-classification/\"\ntraining_image_folder = dataset_folder+\"train/\"\ntrain_df = pd.read_csv(dataset_folder+\"train.csv\")\nlen(train_df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_binary_mask(img):\n    '''\n    Turn the RGB image into grayscale before\n    applying an Otsu threshold to obtain a\n    binary segmentation\n    '''\n    \n    blurred_img = cv2.GaussianBlur(img,(25,25),0)\n    gray_img = cv2.cvtColor(blurred_img, cv2.COLOR_RGBA2GRAY)\n    ret, otsu = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)\n    \n    kernel = np.ones((40,40),np.uint8)\n    closed_mask = cv2.morphologyEx(otsu, cv2.MORPH_CLOSE, kernel)\n    return closed_mask\n\ndef load_RGBY_image(img):\n    '''\n    Load and stack the channels that are stored separately.\n    '''\n    \n    red_image = cv2.imread(img+\"_red.png\", cv2.IMREAD_UNCHANGED)\n    green_image = cv2.imread(img+\"_green.png\", cv2.IMREAD_UNCHANGED)\n    blue_image = cv2.imread(img+\"_blue.png\", cv2.IMREAD_UNCHANGED)\n    yellow_image = cv2.imread(img+\"_yellow.png\", cv2.IMREAD_UNCHANGED)\n\n    stacked_images = np.transpose(np.array([red_image, green_image, blue_image, yellow_image]), (1,2,0))\n    return stacked_images\n\nimg = load_RGBY_image('/kaggle/input/hpa-single-cell-image-classification/train/8061ee18-bbb2-11e8-b2ba-ac1f6b6435d0')\nplt.imshow(img[:, :, :3])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def threshold_img(img):\n    \"\"\"Numpy indexing\"\"\"\n    img_thres = img\n    img_thres[ img < 0.5 ] = 0\n    \n    return img_thres\n\n\nclass CellMaskDataset(Dataset):\n    \n    def __init__(self, stacked_transform=None, mask_transform=None, train=False):\n        \n        self.stacked_transform = stacked_transform\n        self.mask_transform = mask_transform\n        self.train = train\n        \n    def __len__(self):\n        \n        if self.train:\n        \n            return len(train_df[:20000])\n        \n        else:\n            \n            return len(train_df[20001:])\n        \n    def __getitem__(self, idx):\n        \n        if self.train:\n            \n            tmp_df = train_df[:20000]\n            image_path = training_image_folder+tmp_df.iloc[idx].ID\n            stacked_images = load_RGBY_image(image_path)\n            binary_mask = get_binary_mask(stacked_images)\n            #plt.imshow(binary_mask)\n            #plt.show()\n            #print(image_path)\n            stacked_images = stacked_images[:,:,:3]\n            \n            binary_mask = threshold_img(binary_mask)\n            binary_mask = self.mask_transform(image=binary_mask)\n            stacked_images = self.stacked_transform(image=stacked_images)\n            \n        else: \n            \n            tmp_df = train_df[20001:20050]\n            image_path = training_image_folder+tmp_df.iloc[idx].ID\n            stacked_images = load_RGBY_image(image_path)\n            binary_mask = get_binary_mask(stacked_images)\n            stacked_images = stacked_images[:,:,:3]\n            \n            binary_mask = threshold_img(binary_mask)\n            binary_mask = self.mask_transform(image=binary_mask)\n            stacked_images = self.stacked_transform(image=stacked_images)\n        \n        return stacked_images['image'], binary_mask['image'].float()\n    \n    \ndef get_transforms(data):\n    \n    if data == 'mask':\n        \n        return Compose([\n            Resize(64, 64), # For more accurately segmented images: Resize(256, 256)\n            #ToGray(p=1.0),\n            Normalize(\n                mean=[0.0],\n                std=[1.0],\n            ),\n            ToTensorV2(),\n\n        ])\n    \n    else:\n        \n        return Compose([\n            Resize(64, 64), # For more accurately segmented images: Resize(256, 256)\n            Normalize(\n                mean=[0.0, 0.0, 0.0],\n                std=[1.0, 1.0, 1.0],\n            ),\n            ToTensorV2(),\n            \n        ])\n\n\n\ntrain_dataset = CellMaskDataset(stacked_transform=get_transforms('stupid'), mask_transform=get_transforms('mask'), train=True)\ntrain_dataloader = DataLoader(train_dataset,\n                                batch_size=1,\n                                shuffle=True,\n                                num_workers=0)\n\nval_dataset = CellMaskDataset(stacked_transform=get_transforms('stupid'), mask_transform=get_transforms('mask'), train=False)\nval_dataloader = DataLoader(val_dataset,\n                                batch_size=1,\n                                shuffle=False,\n                                num_workers=0)\nfrom torchvision.utils import save_image\n\n#torch.set_printoptions(profile=\"full\")\n#for idx, (image, mask) in enumerate(train_dataloader):\n    #if idx == 22:\n        #save_image(mask, \"real.png\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.nn.modules.loss import _Loss\n\n\nclass SoftDiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(SoftDiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = torch.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\nclass First2D(nn.Module):\n    \n    def __init__(self, in_channels, middle_channels, out_channels, dropout=False):\n        super(First2D, self).__init__()\n        \n        layers = [\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        ]\n        \n        if dropout:\n            \n            layers.append(nn.Dropout2d(p=dropout))\n        \n        self.block = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        \n        x = self.block(x)\n        \n        return x\n    \nclass Upsample2D(nn.Module):\n    \n    def __init__(self, in_channels, middle_channels, out_channels, deconv_channels, dropout=False):\n        super(Upsample2D, self).__init__()\n        \n        layers = [\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(out_channels, deconv_channels, kernel_size=2, stride=2)\n        ]\n        \n        if dropout:\n            \n            layers.append(nn.Dropout2d(p=dropout))\n            \n        self.block = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        \n        x = self.block(x)\n        \n        return x\n    \n    \nclass Downsample2D(nn.Module):\n    \n    def __init__(self, in_channels, middle_channels, out_channels, dropout=False, downsample_kernel=2):\n        super(Downsample2D, self).__init__()\n        \n        layers = [\n            nn.MaxPool2d(kernel_size=downsample_kernel),\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        ]\n        \n        if dropout:\n            \n            layers.append(nn.Dropout2d(p=dropout))\n            \n        self.block = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        \n        x = self.block(x)\n        \n        return x\n    \nclass Center2D(nn.Module):\n    \n    def __init__(self, in_channels, middle_channels, out_channels, deconv_channels, dropout=False):\n        super(Center2D, self).__init__()\n        \n        layers = [\n            nn.MaxPool2d(kernel_size=2),\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(out_channels, deconv_channels, kernel_size=2, stride=2)\n        ]\n        \n        if dropout:\n            \n            layers.append(nn.Dropout2d(p=dropout))\n            \n        self.block = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        \n        x = self.block(x)\n        \n        return x\n    \nclass Last2D(nn.Module):\n    \n    def __init__(self, in_channels, middle_channels, out_channels):\n        super(Last2D, self).__init__()\n        \n        layers = [\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, middle_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(middle_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=1)\n        ]\n        \n        self.block = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        \n        x = self.block(x)\n        \n        return x\n            \n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class UNET(nn.Module):\n    \n    def __init__(self, in_channels, out_channels, conv_depths=[64, 128, 256, 512, 1024]):\n        super(UNET, self).__init__()\n        \n        encoder_layers = []\n        encoder_layers.append(First2D(in_channels, conv_depths[0], conv_depths[0]))\n        encoder_layers.extend([Downsample2D(conv_depths[i], conv_depths[i + 1], conv_depths[i+1]) for i in range(len(conv_depths) - 2)])\n        \n        decoder_layers = []\n        decoder_layers.extend([Upsample2D(2 * conv_depths[i + 1], 2 * conv_depths[i],  2 * conv_depths[i],  conv_depths[i])\n                              for i in reversed(range(len(conv_depths) - 2))])\n        decoder_layers.append(Last2D(conv_depths[1], conv_depths[0], out_channels))\n        \n        self.encoder_block = nn.Sequential(*encoder_layers)\n        self.center_block = Center2D(conv_depths[-2], conv_depths[-1], conv_depths[-1], conv_depths[-2])\n        self.decoder_block = nn.Sequential(*decoder_layers)\n        \n    def forward(self, x, return_all=False):\n        \n        x_enc = [x]\n        for enc_layer in self.encoder_block:\n            \n            x_enc.append(enc_layer(x_enc[-1]))\n            \n        x_dec = [self.center_block(x_enc[-1])]\n        \n        for dec_layer_idx, dec_layer in enumerate(self.decoder_block):\n            \n            x_opposite = x_enc[-1 - dec_layer_idx]\n            x_cat = torch.cat([pad_to_shape(x_dec[-1], x_opposite.shape), x_opposite],\n                             dim=1)\n            x_dec.append(dec_layer(x_cat))\n            \n        if not return_all:\n            \n            return x_dec[-1]\n        \n        else:\n            \n            return x_enc + x_dec\n        \n        \ndef pad_to_shape(current, targ_shp):\n    \n    if len(targ_shp) == 4:\n        \n        pad = (0, targ_shp[3] - current.shape[3], 0, targ_shp[2] - current.shape[2])\n        \n    elif len(targ_shp) == 5:\n        \n        pad = (0, targ_shp[4] - current.shape[4], 0, targ_shp[3] - current.shape[3], 0, targ_shp[2] - current.shape[2])\n        \n    return F.pad(current, pad)\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"u = UNET(3, 1).to(device)\n\nsummary(u, input_size=(3, 64, 64))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\nimport matplotlib.pyplot as plt\nfrom torchvision.utils import save_image\n\nclass Trainer():\n    \n    def __init__(self, net, loss, optimizer, device, scheduler=None):\n        \n        self.net = net\n        self.loss = loss\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        \n        self.device = device\n        self.net.to(self.device)\n        self.loss.to(self.device)\n        \n    def val_epoch(self, dataloader, dice_loss):\n        \n        running_val_loss = 0.0\n        \n        with torch.no_grad():\n            \n            for batch_idx, (x_batch, y_batch) in enumerate(dataloader):\n                \n                x_batch = x_batch.to(self.device)\n                y_batch = y_batch.to(self.device)\n                \n                y_out =  self.net(x_batch)\n                validation_loss = self.loss(y_out, y_batch)\n                dice_loss = dice_loss(y_out, y_batch)\n                running_val_loss += validation_loss\n                \n        return ((running_val_loss/(len(dataloader)), dice_loss))\n        \n    def train_epoch(self, dataloader, dice_loss, epoch):\n        \n        self.net.train()\n        \n        epoch_running_loss = 0.0\n        running_dice_loss = 0.0\n        \n        for batch_idx, (x_batch, y_batch) in enumerate(dataloader):\n            \n            x_batch = x_batch.to(self.device)\n            y_batch = y_batch.to(self.device)\n            \n            self.optimizer.zero_grad()\n            y_out = self.net(x_batch)\n            training_loss = self.loss(y_out, y_batch)\n            train_dice = dice_loss(y_out, y_batch)\n            \n            training_loss.backward()\n            self.optimizer.step()\n            \n            epoch_running_loss += training_loss\n            running_dice_loss += train_dice\n            \n            if batch_idx % 200 == 0 and batch_idx != 0:\n                \n                print(f\"Step - {batch_idx} | Training Loss - {epoch_running_loss/batch_idx} | Dice Loss - {running_dice_loss/batch_idx}\")\n                \n                img1 = torch.sigmoid(y_out) # output is the output tensor of your UNet, the sigmoid will center the range around 0.\n                # Binarize the image\n                threshold = (img1.min() + img1.max()) * 0.5\n                ima = torch.where(img1 > threshold, 1.0, 0.0)\n                \n                #print(ima)\n                \n                save_image(ima, f'BIN_ima_{batch_idx}_{epoch}.png')\n                save_image(y_batch.squeeze(1), f'BIN_ima_groundtruth_{batch_idx}_{epoch}.png')\n                save_image(x_batch.squeeze(1), f'BIN_original_{batch_idx}_{epoch}.png')\n                \n            \n            \n            epoch_running_loss += training_loss.item()\n            \n        return (epoch_running_loss/len(dataloader)), train_dice\n    \n    \n                \n    \n    def train_unet(self, train_loader, val_loader, n_epochs, dice_metric):\n            \n        min_loss = np.inf\n        train_time = time.time()\n        dice_metric = dice_metric.to(device)\n                \n        logs = {}\n        \n        for epoch in range(1, n_epochs+1):\n            \n            train_loss, train_dice = self.train_epoch(train_loader, dice_metric, epoch)\n            \n            #self.scheduler.step()\n            \n            val_loss, val_dice = self.val_epoch(val_loader, dice_metric)\n            \n            logs = {'epoch': epoch,\n                    'time': epoch_end - train_start,\n                    'train_loss': train_loss,\n                    'validation_loss': val_loss,\n                    'train_dice': trian_dice,\n                    'validation_dice': val_dice\n                    }\n            \n            print(\"-\" * 20)\n            print(f\"Epoch - {logs['epoch']} | Time Elapsed - {logs['time']} | Training Loss - {logs['train_loss']} | Train Dice Coeff - {logs['train_dice']}\") \n            print(f\"Validation Loss - {logs['validation_loss']} | Validation Dice - {logs['validation_dice']}\")\n                \n    def predict_dataset(self, dataloader, export_path):\n        \n        with torch.no_grad():\n                \n            for batch_idx, (x_batch) in enumerate(dataloader):\n\n                image_filename = '%s.png' % str(batch_idx + 1).zfill(3)\n\n                x_batch = x_batch.to(self.device)\n                y_out = self.net(x_batch)\n                \n                img1 = torch.sigmoid(y_out) # output is the output tensor of your UNet, the sigmoid will center the range around 0.\n                # Binarize the image\n                threshold = (img1.min() + img1.max()) * 0.5\n                ima = torch.where(img1 > threshold, 1.0, 0.0)\n\n                save_image(ima, os.path.join(export_path, image_filename))\n\n            \n            ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"in_channels = 3\nout_channels = 1\nwidth = 32\ndepth = 6\nconv_depths = [int(width*(2**k)) for k in range(depth)]\n\nunet = UNET(in_channels, out_channels, conv_depths).to(device)\nloss = nn.BCEWithLogitsLoss()\ndice = SoftDiceLoss()\noptimizer = optim.Adam(unet.parameters(), lr=1e-3)\n\ntrainer = Trainer(unet, loss, optimizer, device=device)\n\n#summary(unet, input_size=(3, 256, 256))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.train_unet(train_dataloader, val_dataloader, 5, dice)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%pylab inline\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\n\nprint(\"Original Image\")\nimg = mpimg.imread('./BIN_original_1000_1.png')\nimgplot = plt.imshow(img)\nplt.show()\n\nprint(\"Ground Truth\")\nimg = mpimg.imread('./BIN_ima_groundtruth_1000_1.png')\nimgplot = plt.imshow(img)\nplt.show()\n\nprint(\"Predicted\")\nimg = mpimg.imread('./BIN_ima_1000_1.png')\nimgplot = plt.imshow(img)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}