{"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":"import gc\nimport math\nimport glob\nimport os\nimport torch\nimport numpy as np\nimport torch.nn as nn\nimport torchvision.transforms.functional as TF\nimport torchvision.utils\nimport pytorch_lightning as pl\nimport torchmetrics as tm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom typing import List\nfrom PIL import Image","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:38.87905Z","iopub.execute_input":"2022-06-15T00:10:38.879782Z","iopub.status.idle":"2022-06-15T00:10:38.886607Z","shell.execute_reply.started":"2022-06-15T00:10:38.879734Z","shell.execute_reply":"2022-06-15T00:10:38.885416Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check for GPU or setup TPU\nif not torch.cuda.is_available():\n !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n !python pytorch-xla-env-setup.py --version 1.7 --apt-packages libomp5 libopenblas-dev\n !pip install pytorch-lightning==1.1.5","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-15T00:10:38.959056Z","iopub.execute_input":"2022-06-15T00:10:38.959914Z","iopub.status.idle":"2022-06-15T00:10:38.968213Z","shell.execute_reply.started":"2022-06-15T00:10:38.959879Z","shell.execute_reply":"2022-06-15T00:10:38.967024Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Architecture\n\nThe following image depicts the architecture that will be used for segmentation. We will first define a torch module as the building block for the model, then use a pytorch lightning module to define the final model.\n\n","metadata":{}},{"cell_type":"code","source":"class Block(pl.LightningModule):\n def __init__(self, in_channels, out_channels):\n super().__init__()\n self.conv = nn.Sequential(\n nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias= False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(),\n nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias= False),\n nn.BatchNorm2d(out_channels),\n nn.ReLU()\n )\n def forward(self, x):\n return self.conv(x)\n\nclass UNET(pl.LightningModule):\n def __init__(self, in_channels: int=3, out_channels: int=1, features: List=[64,128,256,512], learning_rate=1.5e-3):\n super().__init__()\n self.learning_rate = learning_rate \n self.down = nn.ModuleList()\n self.up = nn.ModuleList()\n self.pool = nn.MaxPool2d(2, 2)\n for feature in features:\n self.down.append(Block(in_channels, feature))\n in_channels=feature\n for feature in reversed(features):\n self.up.append(\n nn.ConvTranspose2d(feature*2, feature, 2, 2)\n )\n self.up.append(\n Block(feature*2, feature) # x gets concat to 2xchannel\n )\n self.bottleneck = Block(features[-1], features[-1]*2)\n self.final_conv = nn.Conv2d(features[0], out_channels, 1)\n self.loss_fn = nn.BCEWithLogitsLoss()\n \n self.val_num_correct = 0\n self.val_num_pixels = 0\n self.val_dice_score = 0\n self.num_correct = 0\n self.num_pixels = 0\n self.dice_score = 0\n def forward(self, x):\n skip_connections = []\n for down in self.down:\n x = down(x)\n skip_connections.append(x)\n x = self.pool(x)\n x = self.bottleneck(x)\n skip_connections = skip_connections[::-1]\n for idx in range(0, len(self.up), 2):\n x = self.up[idx](x)\n skip_connection = skip_connections[idx//2]\n if x.shape != skip_connection.shape:\n x = TF.resize(x, size=skip_connection.shape[2:])\n concat_skip = torch.cat((skip_connection, x), dim=1) # Concat along channels (b, c, h, w)\n x = self.up[idx+1](concat_skip)\n return self.final_conv(x)\n \n def training_step(self, batch, batch_idx):\n x, y = batch\n pred = self(x)\n loss = self.loss_fn(pred, y)\n pred = torch.sigmoid(pred)\n pred = (pred > 0.5).float()\n self.num_correct += (pred == y).sum()\n self.num_pixels += torch.numel(pred)\n self.dice_score += (2 * (pred * y).sum()) / (\n (pred + y).sum() + 1e-8\n )\n self.log('train_loss', loss, logger = True)\n return {'loss': loss}\n \n def training_epoch_end(self, output):\n train_acc = float(f'{(self.num_correct/self.num_pixels)*100:.2f}')\n self.log('train_acc', train_acc, prog_bar = True, logger = True)\n dice_score = self.dice_score/len(output)\n self.log('train_dice_score', dice_score, prog_bar = True, logger = True)\n self.num_correct, self.num_pixels, self.dice_score = 0,0,0\n \n def validation_step(self, batch, batch_idx):\n x, y = batch\n pred = self(x)\n loss = self.loss_fn(pred, y)\n pred = torch.sigmoid(pred)\n pred = (pred > 0.5).float()\n self.val_num_correct += (pred == y).sum()\n self.val_num_pixels += torch.numel(pred)\n self.val_dice_score += (2 * (pred * y).sum()) / (\n (pred + y).sum() + 1e-8\n )\n self.log('val_loss', loss, prog_bar = True, logger = True)\n return {'loss': loss}\n \n def validation_epoch_end(self, output):\n val_acc = float(f'{(self.val_num_correct/self.val_num_pixels)*100:.2f}')\n self.log('val_acc', val_acc, prog_bar = True, logger = True)\n dice_score = self.val_dice_score/len(output)\n self.log('val_dice_score', dice_score, prog_bar = True, logger = True)\n self.val_num_correct, self.val_num_pixels, self.val_dice_score = 0,0,0\n \n def configure_optimizers(self):\n optimizer = torch.optim.AdamW(params = self.parameters(), lr = self.learning_rate, weight_decay=0.3)\n return optimizer","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-15T00:10:39.074425Z","iopub.execute_input":"2022-06-15T00:10:39.075059Z","iopub.status.idle":"2022-06-15T00:10:39.105093Z","shell.execute_reply.started":"2022-06-15T00:10:39.075022Z","shell.execute_reply":"2022-06-15T00:10:39.104171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can see how many parameters we have in our model using the following command. Note that the order shown does not reflect how a training example flows through the architecture. ","metadata":{}},{"cell_type":"code","source":"pl.utilities.model_summary.summarize(UNET(),-1)","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:39.183976Z","iopub.execute_input":"2022-06-15T00:10:39.186706Z","iopub.status.idle":"2022-06-15T00:10:39.478585Z","shell.execute_reply.started":"2022-06-15T00:10:39.186664Z","shell.execute_reply":"2022-06-15T00:10:39.4776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data\nNow that we have our model, we define a torch Dataset to retrieve and apply our transformations to our data. We define the three necessary methods and move on to defining our LightningDataModule, which will split our Dataset into train/val splits and prepare the appropriate Dataloader when called by the Trainer object later on when we start training.","metadata":{}},{"cell_type":"code","source":"class SegmentationDataset(torch.utils.data.Dataset):\n def __init__(self, image_path, mask_path, transforms):\n self.images = glob.glob(os.path.join(image_path, '*.jpg'))\n self.image_path = image_path\n self.mask_path = mask_path\n self.transforms = transforms\n\n def __len__(self):\n return len(self.images)\n \n def __getitem__(self, idx):\n img = np.array(Image.open(self.images[idx]).convert('RGB'))\n mask = np.array(Image.open(os.path.join(self.mask_path, os.path.basename(self.images[idx]).replace('.jpg', '.png')))) \n mask[mask == 255.0] = 1.0 \n augmentations = self.transforms(image=img, mask=mask)\n image = augmentations[\"image\"]\n mask = augmentations[\"mask\"]\n mask = torch.unsqueeze(mask, 0)\n mask = mask.type(torch.float32)\n return image, mask","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:39.480537Z","iopub.execute_input":"2022-06-15T00:10:39.48097Z","iopub.status.idle":"2022-06-15T00:10:39.491225Z","shell.execute_reply.started":"2022-06-15T00:10:39.480912Z","shell.execute_reply":"2022-06-15T00:10:39.490229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SegmentationDataModule(pl.LightningDataModule):\n \n def __init__(self, image_path, mask_path, transform, train_size=0.90, batch_size: int = 9):\n super().__init__()\n self.image_path = image_path\n self.mask_path = mask_path\n self.batch_size = batch_size\n self.transform = transform\n self.train_size = train_size\n \n def setup(self, stage = None):\n if stage in (None, 'fit'):\n ds = SegmentationDataset(self.image_path, self.mask_path, self.transform)\n train_size = math.floor(len(ds)*self.train_size)\n val_size = len(ds)-train_size\n train_ds, val_ds = torch.utils.data.random_split(ds, [train_size, val_size])\n self.train_dataset = train_ds\n self.val_dataset = val_ds\n \n def train_dataloader(self):\n return torch.utils.data.DataLoader(self.train_dataset, self.batch_size, num_workers=2, shuffle = True, persistent_workers=True)\n\n def val_dataloader(self):\n return torch.utils.data.DataLoader(self.val_dataset, self.batch_size, num_workers=2, persistent_workers=True)\n \n def test_dataloader(self):\n return torch.utils.data.DataLoader(self.val_dataset, self.batch_size)","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:39.492677Z","iopub.execute_input":"2022-06-15T00:10:39.493639Z","iopub.status.idle":"2022-06-15T00:10:39.507112Z","shell.execute_reply.started":"2022-06-15T00:10:39.493536Z","shell.execute_reply":"2022-06-15T00:10:39.506212Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training\n\nOnce our model, optimizer+loss, and LightningDataModule is defined, we can use Trainer to run our training and validation loops. We can also use callbacks provided by lightning (check out [Bolts](https://lightning-bolts.readthedocs.io/en/latest/) for advanced callbacks, such as for sparsification) to create checkpoints or for early stopping. Trainer also takes care of multi-GPU and TPU training. \n\nWe also use albumentations to apply the **same** data augmentations on the image **and** mask. I don't believe the torchvision transforms allows you to do the same, but correct me if I'm wrong. ","metadata":{}},{"cell_type":"code","source":"torch.cuda.empty_cache()\ntransform = A.Compose(\n [\n A.Resize(height=360, width=480),\n A.Rotate(limit=45, p=0.7),\n A.HorizontalFlip(p=0.5),\n A.VerticalFlip(p=0.3),\n A.Normalize(\n mean=[0.0, 0.0, 0.0],\n std=[1.0, 1.0, 1.0],\n max_pixel_value=255.0,\n ),\n A.pytorch.ToTensorV2(),\n ]\n)\nds = SegmentationDataModule('../input/carvana-image-masking-png/train_images', '../input/carvana-image-masking-png/train_masks', transform=transform)\nmodel = UNET()\nif not os.path.isdir('./unet_3'): os.mkdir('./unet_3')\ncheckpointCallback = pl.callbacks.ModelCheckpoint(dirpath=\"./unet_3\", \n save_top_k=1, \n monitor=\"val_loss\",\n filename='{epoch}-{val_loss:.5f}',\n mode='min')\nif torch.cuda.is_available():\n trainer = pl.Trainer(max_epochs=8, accelerator='gpu', gpus=1, \n callbacks=[checkpointCallback], profiler='simple',\n auto_lr_find=True)\nelse:\n trainer = pl.Trainer(max_epochs=7, accelerator='tpu', tpu_cores=8, \n callbacks=[checkpointCallback], profiler='simple', \n auto_lr_find=True)\n\ntrainer.fit(model, datamodule=ds)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-06-15T00:10:39.509388Z","iopub.execute_input":"2022-06-15T00:10:39.510287Z","iopub.status.idle":"2022-06-15T00:10:57.406874Z","shell.execute_reply.started":"2022-06-15T00:10:39.51023Z","shell.execute_reply":"2022-06-15T00:10:57.405516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utility Functions\n\nHere we define three utility functions:\n 1. **save_images**(model, loader, folder, device)\n - This is used to save predictions and their corresponding masks from a specified model across a specified dataloader.\n - Device can be set to cpu if not running on gpu\n 2. **get_concat_v**(im1, im2)\n - This is used to later concatenate the prediction and mask images we create using save_image.\n - im1 & im2 are PIL.Image objects\n 3. **merge_photos**(src_folder, dst_folder, remove_single)\n - This is used to read and concatenate the prediction and mask images using the previously defined get_concat_v().\n - remove_single can be set to False to save the unmerged images","metadata":{}},{"cell_type":"code","source":"def save_images(model, loader, folder='val_img', device='cuda'):\n model.eval()\n if not os.path.isdir(folder):\n os.mkdir(folder)\n model.to(device=device)\n for idx, (x, y) in enumerate(loader):\n x = x.to(device=device)\n y = y.to(device=device)\n with torch.no_grad():\n preds = torch.sigmoid(model(x).cuda())\n preds = (preds > 0.5).float()\n torchvision.utils.save_image(\n preds, f\"{folder}/pred_{idx}.png\"\n )\n torchvision.utils.save_image(y, f\"{folder}/mask_{idx}.png\") \n\ndef get_concat_v(im1, im2):\n dst = Image.new('RGB', (im1.width, im1.height + im2.height))\n dst.paste(im1, (0, 0))\n dst.paste(im2, (0, im1.height))\n return dst\n\ndef merge_photos(src_folder: str='./val_img', dst_folder: str='./merged_val_img', remove_single: bool=True):\n files = glob.glob(src_folder+'/*.png')\n if not os.path.isdir(dst_folder):\n os.mkdir(dst_folder)\n for i in range(int(len(files)/2)):\n pred_img = Image.open(f'{src_folder}/pred_{i}.png')\n mask_img = Image.open(f'{src_folder}/mask_{i}.png')\n get_concat_v(pred_img, mask_img).save(f'{dst_folder}/merged_pred_mask_{i}.png')\n if remove_single:\n os.remove(f'./val_img/pred_{i}.png')\n os.remove(f'./val_img/mask_{i}.png')","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:57.408658Z","iopub.execute_input":"2022-06-15T00:10:57.409192Z","iopub.status.idle":"2022-06-15T00:10:57.420929Z","shell.execute_reply.started":"2022-06-15T00:10:57.409147Z","shell.execute_reply":"2022-06-15T00:10:57.420225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds.setup()\nsave_images(model, ds.val_dataloader())\nmerge_photos()\n# TODO: Implement feature for sampling random img/label pairs for model prediction/eval\n# - Should show output and mask as images, pref side by side or across batches\n# - Conduct additional testing using the competition or data from the internet","metadata":{"execution":{"iopub.status.busy":"2022-06-15T00:10:57.422114Z","iopub.execute_input":"2022-06-15T00:10:57.423057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!zip -r lightning_logs.zip ./lightning_logs\n!zip -r merged_val_img.zip ./merged_val_img\n!zip -r unet.zip ./unet_3","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Since Kaggle doesn't work with Tensorboard, download and load the logs in Google Collab to analyze the training and validation metrics\n\n\n\n","metadata":{}}]}