{"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":"# 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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config\n","metadata":{}},{"cell_type":"code","source":"import torch","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Hyperparameters etc.\nLEARNING_RATE = 1e-4\nWEIGHT_DECAY = 2e-5\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 32\nTEST_BATCH_SIZE = 128\nNUM_EPOCHS = 50\nNUM_WORKERS = 2\nIMAGE_HEIGHT = 160\nIMAGE_WIDTH = 240\nPIN_MEMORY = True\nLOAD_MODEL = False","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils\n","metadata":{}},{"cell_type":"code","source":"import glob\nimport os\nfrom pathlib import Path\n\nimport torch\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, save_dir='runs/', file_name='my_checkpoint.pt'):\n    print('=> Saving checkpoint')\n    save_path = os.path.join(save_dir, file_name)\n    checkpoint = {\n        'state_dict': model.state_dict(),\n        'optimizer': optimizer.state_dict()\n    }\n    torch.save(checkpoint, save_path)\n\n\ndef load_checkpoint(model, optimizer, lr, load_dir='runs/', file_name='my_checkpoint.pt'):\n    print('=> Loading checkpoint')\n    load_path = os.path.join(load_dir, file_name)\n    checkpoint = torch.load(load_path, map_location=DEVICE)\n    model.load_state_dict(checkpoint['state_dict'])\n    optimizer.load_state_dict(checkpoint['optimizer'])\n\n    for params_group in optimizer.params_groups:\n        params_group['lr'] = lr\n        \n        \ndef check_accuracy(loader, model, loss_fn, device='cuda'):\n    num_correct = 0\n    num_pixels = 0\n    dice_score = 0\n    loss = 0\n\n    model.eval()\n\n    with torch.no_grad():\n        for x, y in loader:\n            x = x.to(device)\n            y = y.to(device).unsqueeze(1)\n            preds = model(x)\n            loss += loss_fn(preds, y).item() * x.size(0)\n            \n            preds = torch.sigmoid(preds)\n            preds = (preds >= .5).float()\n            num_correct += (preds == y).sum()\n            num_pixels += torch.numel(preds)\n            dice_score += (2 * (preds * y).sum()) / ((preds + y).sum() + 1e-8)\n            \n    loss = loss / len(loader.dataset)\n\n    print(\n        f'Got {num_correct}/{num_pixels} with acc {num_correct/num_pixels*100:.2f}\\t'\n        f'Dice score: {dice_score/len(loader)}\\t'\n        f'Loss: {loss}'\n    )\n    model.train()\n\n    return num_correct/num_pixels*100, dice_score/len(loader), loss\n\n\ndef save_predictions_as_imgs(\n    loader, model, folder='save_images/', device='cuda', num=5,\n):\n    model.eval()\n    for idx, (x, y) in enumerate(loader):\n        x = x.to(device)\n        with torch.no_grad():\n            preds = torch.sigmoid(model(x))\n            preds = (preds > .5).float()\n        torchvision.utils.save_image(\n            preds, f'{folder}pred_{idx}.png'\n        )\n        torchvision.utils.save_image(y.unsqueeze(1), f'{folder}{idx}.png')\n        \n        if idx > num:\n            break\n\n    model.train()\n    \n    \ndef visualize(display_list):\n    plt.figure(figsize=(15, 20))\n    title = ['Input image', 'True mask', 'Predicted mask']\n    for i in range(len(display_list)):\n        plt.subplot(1, len(display_list), i+1)\n        plt.title(title[1])\n        plt.imshow(display_list[i])\n        plt.axis('off')\n    plt.show()\n    \n    \ndef show_predictions(image, mask, device='cuda'):\n    model.eval()\n    \n    image = image.unsqueeze(0).to(device)\n    pred_mask = model(image).squeeze()\n    \n    image = image.squeeze(0).cpu().numpy().transpose(1, 2, 0)\n    mask = mask.numpy()\n    pred_mask = pred_mask.detach().cpu().numpy()\n    \n    visualize([image, mask, pred_mask])\n    \n    model.train()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset\n\n","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport os\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom zipfile import ZipFile\nfrom PIL import Image\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import train_test_split\n\nimport numpy as np\nimport torch\nfrom torch.utils.data import DataLoader, Dataset","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train data\ntrain_zip = '../input/carvana-image-masking-challenge/train.zip'\n\nwith ZipFile(train_zip, 'r') as zip_:\n    zip_.extractall()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train mask data\ntrain_masks_zip = '../input/carvana-image-masking-challenge/train_masks.zip'\n\nwith ZipFile(train_masks_zip, 'r') as zip_:\n    zip_.extractall()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train_csv data\ntrain_csv_zip = '../input/carvana-image-masking-challenge/train_masks.csv.zip'\n\nwith ZipFile(train_csv_zip, 'r') as zip_:\n    zip_.extractall()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"car_ids = []\npaths = []\nfor dir_name, _, file_names in os.walk('./train'):\n    for file_name in file_names:\n        path = os.path.join(dir_name, file_name)\n        paths.append(path)\n        \n        car_id = file_name.split('.')[0]\n        car_ids.append(car_id)\n        \nd = {'id':car_ids, 'car_path':paths}\ndf = pd.DataFrame(data = d)\ndf = df.set_index('id')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"car_ids = []\npaths = []\nfor dir_name, _, file_names in os.walk('./train_masks'):\n    for file_name in file_names:\n        path = os.path.join(dir_name, file_name)\n        paths.append(path)\n        \n        car_id = file_name.split('_mask')[0]\n        car_ids.append(car_id)\n        \nd = {'id':car_ids, 'mask_path':paths}\nmask_df = pd.DataFrame(data = d)\nmask_df = mask_df.set_index('id')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df['mask_path'] = mask_df['mask_path']\n\ndf","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CarvanaDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.transform = transform\n        self.images = df['car_path'].tolist()\n        self.masks = df['mask_path'].tolist()\n        assert len(self.images) == len(self.masks), 'number of images data not equal to masks data'\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        img_path = self.images[idx]\n        mask_path = self.masks[idx]\n        image = np.array(Image.open(img_path).convert('RGB'))\n        mask = np.array(Image.open(mask_path).convert('L'), dtype=np.float32)\n        mask[mask == 255.] = 1.0\n\n        if self.transform is not None:\n            augmentations = self.transform(image = image, mask = mask)\n            image = augmentations['image']\n            mask = augmentations['mask']\n\n        return image, mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Augmentations\n\ntrain_transforms = A.Compose(\n    [\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n        A.Rotate(limit=35, p=1.0),\n        A.HorizontalFlip(p=.1),\n        A.VerticalFlip(p=.5),\n        A.Normalize(\n            mean=[0., 0., 0.],\n            std=[1., 1., 1.],\n            max_pixel_value=255.,\n        ),\n        ToTensorV2(),\n    ]\n)\n\nval_transforms = A.Compose(\n    [\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n        A.Normalize(\n            mean=[0., 0., 0.],\n            std=[1., 1., 1.],\n            max_pixel_value=255.,\n        ),\n        ToTensorV2(),\n    ]\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df, valid_df = train_test_split(df, random_state=2103, test_size=.2)\n\ntrain_dataset = CarvanaDataset(train_df, train_transforms)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY)\n\nvalid_dataset = CarvanaDataset(valid_df, val_transforms)\nvalid_loader = DataLoader(valid_dataset, batch_size=TEST_BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check data\ndataiter = iter(train_loader)\nimages, masks = dataiter.next()\n\nprint(images.shape)\nprint(masks.shape)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualize some groundtruth\nrand_idx = np.random.randint(len(train_dataset), size=5)\n\nfig, axes = plt.subplots(5, 2, figsize=(15, 20))\n\n# Plot the images\nfor i, idx in enumerate(rand_idx):\n    img, mask = train_dataset[idx]\n    img = img.numpy()\n    mask = mask.numpy()\n    img = img.transpose(1, 2, 0)\n    \n    ax_img = axes[i][0]\n    ax_mask = axes[i][1]\n    \n    ax_img.imshow(img)\n    ax_img.get_xaxis().set_visible(False)\n    ax_img.get_yaxis().set_visible(False)\n    ax_img.set_title('original')\n    \n    ax_mask.imshow(mask)\n    ax_mask.get_xaxis().set_visible(False)\n    ax_mask.get_yaxis().set_visible(False)\n    ax_mask.set_title('mask')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torchvision.transforms.functional as TF\nfrom torch.nn.init import kaiming_uniform_, xavier_uniform_","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DoubleConv, self).__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.LeakyReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.LeakyReLU(inplace=True),\n        )\n        kaiming_uniform_(self.conv[0].weight, nonlinearity='relu')\n        kaiming_uniform_(self.conv[3].weight, nonlinearity='relu')\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass UNET(nn.Module):\n    def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):\n        super(UNET, self).__init__()\n        self.ups = nn.ModuleList()\n        self.downs = nn.ModuleList()\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        # Down part of Unet\n        for feature in features:\n            self.downs.append(DoubleConv(in_channels, feature))\n            in_channels = feature\n\n        # Up part of Unet\n        for feature in reversed(features):\n            self.ups.append(\n                nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2),\n            )\n            self.ups.append(DoubleConv(feature*2, feature))\n\n        self.bottleneck = DoubleConv(features[-1], features[-1]*2)\n        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)\n\n    def forward(self, x):\n        skip_connections = []\n\n        for down in self.downs:\n            x = down(x)\n            skip_connections.append(x)\n            x = self.pool(x)\n        \n        x = self.bottleneck(x)\n        skip_connections = skip_connections[::-1]\n\n        for idx in range(0, len(self.ups), 2):\n            x = self.ups[idx](x)\n            skip_connection = skip_connections[idx//2]\n            \n            if x.shape != skip_connection.shape:\n                x = TF.resize(x, skip_connection.shape[2:])\n\n            concat_skip = torch.cat((skip_connection, x), dim=1)\n            x = self.ups[idx + 1](concat_skip)\n\n        return self.final_conv(x)\n    \n\nmodel = UNET(3, 1)\nmodel","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"from tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.optim.lr_scheduler as lr_scheduler","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(loader, model, optimizer, scheduler, loss_fn, scaler):\n    loop = tqdm(loader)\n\n    for _, (data, targets) in enumerate(loop):\n        data = data.float().to(DEVICE)\n        targets = targets.float().unsqueeze(1).to(DEVICE)\n\n        # forward\n        with torch.cuda.amp.autocast():\n            predictions = model(data)\n            loss = loss_fn(predictions, targets)\n\n        # backward\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # update tqdm loop\n        loop.set_postfix(loss=loss.item())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir 'predict_with_best_dice_score/'\n!mkdir 'predict_with_best_loss/'\n!mkdir 'runs/'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_fn = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=.75, patience=5, verbose=5, eps=1e-6)\n\nscaler = torch.cuda.amp.GradScaler()\n\ntrain_accuracy = []\ntrain_dice_score = []\ntrain_loss = []\nvalid_accuracy = []\nvalid_dice_score = []\nvalid_loss = []\n\nbest_dice_score = 0\nbest_loss = np.Inf\n\nmodel.to(DEVICE)\n\nfor epoch in range(NUM_EPOCHS):\n    print(f'epoch: {epoch}/{NUM_EPOCHS}')\n    train_fn(train_loader, model, optimizer, scheduler, loss_fn, scaler)\n    \n    # Check accuracy and loss\n    print('TRAIN:\\t', end='')\n    train_acc, train_d_score, train_l = check_accuracy(train_loader, model, loss_fn, DEVICE)\n    train_accuracy.append(train_acc)\n    train_dice_score.append(train_d_score)\n    train_loss.append(train_l)\n    \n    print('VALID:\\t', end='')\n    valid_acc, valid_d_score, valid_l = check_accuracy(valid_loader, model, loss_fn, DEVICE)\n    valid_accuracy.append(valid_acc)\n    valid_dice_score.append(valid_d_score)\n    valid_loss.append(valid_l)\n    \n    # Updates Learning Rate wrt the Average Validation Loss\n    scheduler.step(valid_loss[-1])\n    \n    if best_dice_score < valid_d_score:\n        print(f'\\tSave best dice score: {best_dice_score} --> {valid_d_score}')\n        best_dice_score = valid_d_score\n        save_checkpoint(model, optimizer, file_name='best_dice_score.pt')\n        save_predictions_as_imgs(valid_loader, model, folder='predict_with_best_dice_score/', device=DEVICE)\n        \n    if best_loss > valid_loss[-1]:\n        print(f'\\tSave best loss: {best_loss} --> {valid_loss}')\n        best_loss = valid_loss[-1]\n        save_checkpoint(model, optimizer, file_name='best_loss.pt')\n        save_predictions_as_imgs(valid_loader, model, folder='predict_with_best_loss/', device=DEVICE)\n    \n    show_predictions(valid_dataset[0][0], valid_dataset[0][1], device=DEVICE)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# test_img = np.array(Image.open('./train/e8f607100c1f_14.jpg').convert('RGB'))\n# print(test_img.shape)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !unzip ../input/carvana-image-masking-challenge/sample_submission.csv.zip","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# sub = pd.read_csv('./sample_submission.csv')\n\n# sub","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}