{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":6927,"databundleVersionId":45059,"sourceType":"competition"}],"dockerImageVersionId":30646,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nprint(os.listdir(\"../input/carvana-image-masking-challenge/\"))\n\nimport zipfile\nimport shutil\n\nDATASET_DIR = '../input/carvana-image-masking-challenge/'\nWORKING_DIR = '/kaggle/working/'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-02-21T08:25:42.400047Z","iopub.execute_input":"2024-02-21T08:25:42.400405Z","iopub.status.idle":"2024-02-21T08:25:42.406397Z","shell.execute_reply.started":"2024-02-21T08:25:42.400376Z","shell.execute_reply":"2024-02-21T08:25:42.405539Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if len(os.listdir(WORKING_DIR)) <= 1:\n\n    with zipfile.ZipFile(DATASET_DIR + 'train.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n\n    with zipfile.ZipFile(DATASET_DIR + 'train_masks.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n    \n    print(\n        len(os.listdir(WORKING_DIR + 'train')),\n        len(os.listdir(WORKING_DIR + 'train_masks'))\n    )\n\n    train_dir = WORKING_DIR + 'train/'\n    val_dir = WORKING_DIR + 'val/'\n    os.mkdir(val_dir)\n    for file in sorted(os.listdir(train_dir))[4600:]:\n      shutil.move(train_dir + file, val_dir)\n\n    masks_dir = WORKING_DIR + 'train_masks/'\n    val_masks_dir = WORKING_DIR + 'val_masks/'\n    os.mkdir(val_masks_dir)\n    for file in sorted(os.listdir(masks_dir))[4600:]:\n      shutil.move(masks_dir + file, val_masks_dir)\n    \n    os.mkdir(WORKING_DIR + 'saved_images')","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:25:45.114729Z","iopub.execute_input":"2024-02-21T08:25:45.115344Z","iopub.status.idle":"2024-02-21T08:25:53.220964Z","shell.execute_reply.started":"2024-02-21T08:25:45.115312Z","shell.execute_reply":"2024-02-21T08:25:53.220041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sorted(os.listdir(WORKING_DIR + 'train'))[:10]","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:26:57.155583Z","iopub.execute_input":"2024-02-21T08:26:57.156273Z","iopub.status.idle":"2024-02-21T08:26:57.166487Z","shell.execute_reply.started":"2024-02-21T08:26:57.156245Z","shell.execute_reply":"2024-02-21T08:26:57.165562Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sorted(os.listdir(WORKING_DIR + 'train_masks'))[:10]","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:26:58.307223Z","iopub.execute_input":"2024-02-21T08:26:58.308062Z","iopub.status.idle":"2024-02-21T08:26:58.318957Z","shell.execute_reply.started":"2024-02-21T08:26:58.308032Z","shell.execute_reply":"2024-02-21T08:26:58.318067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom PIL import Image\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nclass CarvanaDataset(Dataset):\n    def __init__(self, image_dir, mask_dir, transform=None):\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n        self.images = os.listdir(image_dir)\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_path = os.path.join(self.image_dir, self.images[index])\n        mask_path = os.path.join(self.mask_dir, self.images[index].replace('.jpg', '_mask.gif'))\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.0] = 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":{"execution":{"iopub.status.busy":"2024-02-21T08:27:13.301527Z","iopub.execute_input":"2024-02-21T08:27:13.301874Z","iopub.status.idle":"2024-02-21T08:27:14.948002Z","shell.execute_reply.started":"2024-02-21T08:27:13.301847Z","shell.execute_reply":"2024-02-21T08:27:14.947139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torchvision.transforms.functional as TF\n\nclass 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, kernel_size=3, stride=1, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:29:16.285146Z","iopub.execute_input":"2024-02-21T08:29:16.285541Z","iopub.status.idle":"2024-02-21T08:29:16.293304Z","shell.execute_reply.started":"2024-02-21T08:29:16.285508Z","shell.execute_reply":"2024-02-21T08:29:16.292293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class 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.downs = nn.ModuleList()\n        self.ups = 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(\n                    feature*2, feature, kernel_size=2, stride=2,\n                )\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], 1, 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, size=skip_connection.shape[2:], antialias=None)\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)","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:29:19.083683Z","iopub.execute_input":"2024-02-21T08:29:19.084591Z","iopub.status.idle":"2024-02-21T08:29:19.100542Z","shell.execute_reply.started":"2024-02-21T08:29:19.084551Z","shell.execute_reply":"2024-02-21T08:29:19.099579Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test():\n    x = torch.randn((3, 1, 161, 161))\n    model = UNET(1, 1)\n    preds = model(x)\n    print(preds.shape)\n    print(x.shape)\n\n    assert preds.shape == x.shape","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:29:28.006053Z","iopub.execute_input":"2024-02-21T08:29:28.006430Z","iopub.status.idle":"2024-02-21T08:29:28.011726Z","shell.execute_reply.started":"2024-02-21T08:29:28.006400Z","shell.execute_reply":"2024-02-21T08:29:28.010689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test()","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:29:28.437401Z","iopub.execute_input":"2024-02-21T08:29:28.437769Z","iopub.status.idle":"2024-02-21T08:29:29.639695Z","shell.execute_reply.started":"2024-02-21T08:29:28.437737Z","shell.execute_reply":"2024-02-21T08:29:29.638751Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torchvision\nfrom torch.utils.data import DataLoader","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:30:20.910343Z","iopub.execute_input":"2024-02-21T08:30:20.910714Z","iopub.status.idle":"2024-02-21T08:30:20.917046Z","shell.execute_reply.started":"2024-02-21T08:30:20.910685Z","shell.execute_reply":"2024-02-21T08:30:20.916256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_checkpoint(state, filename='my_checkpoint.pth.tar'):\n    print('=> Saving checkpoint')\n    torch.save(state, filename)\n\ndef load_checkpoint(checkpoint, model):\n    print('=> Loading checkpoint')\n    model.load_state_dict(checkpoint['state_dict'])\n\ndef get_loaders(\n    train_dir, train_maskdir, val_dir, val_maskdir, batch_size, train_transform, val_transform, num_workers=4, pin_memory=True,):\n    train_ds = CarvanaDataset(\n        image_dir=train_dir,\n        mask_dir=train_maskdir,\n        transform=train_transform\n    )\n\n    train_loader = DataLoader(\n        train_ds,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        pin_memory=pin_memory,\n        shuffle=True\n    )\n\n    val_ds = CarvanaDataset(\n        image_dir=val_dir,\n        mask_dir=val_maskdir,\n        transform=val_transform\n    )\n    \n    val_loader = DataLoader(\n        val_ds,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        pin_memory=pin_memory,\n        shuffle=False\n    )\n\n    return train_loader, val_loader\n\ndef check_accuracy(loader, model, device='cuda'):\n    num_correct = 0\n    num_pixels = 0\n    dice_score = 0\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 = torch.sigmoid(model(x))\n            preds = (preds > 0.5).float()\n            num_correct += (preds == y).sum()\n            num_pixels += torch.numel(preds)\n            dice_score += (2 * (preds * y).sum()) / (\n                (preds + y).sum() + 1e-8\n            )\n\n    print(\n        f'Got {num_correct}/{num_pixels} with acc {num_correct/num_pixels*100:.2f}'\n    )\n    print(f'Dice score: {dice_score/len(loader)}')\n    model.train()\n\ndef save_predictions_as_imgs(\n        loader, model, folder=WORKING_DIR+'saved_images', device='cuda'\n):\n    model.eval()\n    for idx, (x, y) in enumerate(loader):\n        x = x.to(device=device)\n        with torch.no_grad():\n            preds = torch.sigmoid(model(x))\n            preds = (preds > 0.5).float()\n        torchvision.utils.save_image(\n            preds, f'{folder}/pred_{idx}.png'\n        )\n        torchvision.utils.save_image(\n            y.unsqueeze(1), f'{folder}/truth_{idx}.png'\n        )","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:30:31.690837Z","iopub.execute_input":"2024-02-21T08:30:31.691206Z","iopub.status.idle":"2024-02-21T08:30:31.705982Z","shell.execute_reply.started":"2024-02-21T08:30:31.691176Z","shell.execute_reply":"2024-02-21T08:30:31.704968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm import tqdm\n\n# Hyperparameters etc.\nLEARNING_RATE = 1e-4\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 16\nNUM_EPOCHS = 50\nNUM_WORKERS = 2\nIMAGE_HEIGHT = 320\nIMAGE_WIDTH = 480\nPIN_MEMORY = True\nLOAD_MODEL = False\nTRAIN_IMG_DIR = '/kaggle/working/train/'\nTRAIN_MASK_DIR = '/kaggle/working/train_masks/'\nVAL_IMG_DIR = '/kaggle/working/val/'\nVAL_MASK_DIR = '/kaggle/working/val_masks/'","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:30:47.328787Z","iopub.execute_input":"2024-02-21T08:30:47.329155Z","iopub.status.idle":"2024-02-21T08:30:47.335667Z","shell.execute_reply.started":"2024-02-21T08:30:47.329127Z","shell.execute_reply":"2024-02-21T08:30:47.334621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(loader, model, optimizer, loss_fn, scaler):\n    loop = tqdm(loader)\n\n    for batch_idx, (data, targets) in enumerate(loop):\n        data = data.to(device=DEVICE)\n        targets = targets.float().unsqueeze(1).to(device=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        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        # update tqdm loop\n        loop.set_postfix(loss=loss.item())","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:31:12.754769Z","iopub.execute_input":"2024-02-21T08:31:12.755655Z","iopub.status.idle":"2024-02-21T08:31:12.762183Z","shell.execute_reply.started":"2024-02-21T08:31:12.755620Z","shell.execute_reply":"2024-02-21T08:31:12.761219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ntrain_transform = A.Compose(\n    [\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n        A.Rotate(limit=35, p=1.0),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.1),\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        ToTensorV2(),\n    ]\n)\n\nval_transform = A.Compose(\n    [\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\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        ToTensorV2(), \n    ]\n)\n\nmodel = UNET(in_channels=3, out_channels=1).to(DEVICE)\nloss_fn = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n\ntrain_loader, val_loader = get_loaders(\n    TRAIN_IMG_DIR,\n    TRAIN_MASK_DIR,\n    VAL_IMG_DIR,\n    VAL_MASK_DIR,\n    BATCH_SIZE,\n    train_transform,\n    val_transform,\n    NUM_WORKERS,\n    PIN_MEMORY\n)\n\nif LOAD_MODEL:\n    load_checkpoint(torch.load('my_checkpoint.pth.tar'), model)\n    check_accuracy(val_loader, model, device=DEVICE)\n\nscaler = torch.cuda.amp.GradScaler()\n\nfor epoch in range(NUM_EPOCHS):\n    train_fn(train_loader, model, optimizer, loss_fn, scaler)\n\n    # save model\n    checkpoint = {\n        'state_dict': model.state_dict(),\n        'optimizer': optimizer.state_dict(),\n    }\n    save_checkpoint(checkpoint)\n\n    # check accuracy\n    check_accuracy(val_loader, model, device=DEVICE)\n\n    # print some examples to a folder\n    save_predictions_as_imgs(\n        val_loader, model, folder=WORKING_DIR+'saved_images', device=DEVICE\n    )","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:31:33.662990Z","iopub.execute_input":"2024-02-21T08:31:33.663364Z","iopub.status.idle":"2024-02-21T08:31:39.833851Z","shell.execute_reply.started":"2024-02-21T08:31:33.663334Z","shell.execute_reply":"2024-02-21T08:31:39.832639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport time\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[:-1:2]\n    \n    return ' '.join(str(x) for x in runs)\n","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:33:21.093809Z","iopub.execute_input":"2024-02-21T08:33:21.094197Z","iopub.status.idle":"2024-02-21T08:33:21.417570Z","shell.execute_reply.started":"2024-02-21T08:33:21.094158Z","shell.execute_reply":"2024-02-21T08:33:21.416620Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 'test' not in os.listdir(WORKING_DIR):\n    if os.path.isfile(WORKING_DIR + 'submission.csv'):\n        os.remove(WORKING_DIR + 'submission.csv')\n    with zipfile.ZipFile(DATASET_DIR + 'test.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\nif 'test_images' not in os.listdir(WORKING_DIR):\n    os.mkdir(WORKING_DIR + 'test_images')","metadata":{"execution":{"iopub.status.busy":"2024-02-21T08:33:31.065119Z","iopub.execute_input":"2024-02-21T08:33:31.065858Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TEST_DIR = WORKING_DIR + 'test'\nTHRESHOLD = 0.5\n\n# Dataset\nclass CarvanaTestDataset(Dataset):\n    def __init__(self, image_dir, transform=None):\n        self.image_dir = image_dir\n        self.transform = transform\n        self.images = sorted(os.listdir(image_dir))\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_name = self.images[index]\n        img_path = os.path.join(self.image_dir, self.images[index])\n        image = np.array(Image.open(img_path).convert('RGB'))\n\n        if self.transform is not None:\n            augmentations = self.transform(image=image)\n            image = augmentations['image']\n\n        return img_name, image\n\n\ntest_transform = A.Compose(\n    [\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\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        ToTensorV2(), \n    ]\n)\n   \ntest_set = CarvanaTestDataset(\n    image_dir=TEST_DIR,\n    transform=test_transform\n)    \n\n\ntest_loader = DataLoader(\n    test_set, batch_size=BATCH_SIZE, shuffle=False\n)\n    \n# Model\ncheckpoint = torch.load(WORKING_DIR + 'my_checkpoint.pth.tar')\n\nmodel = UNET(in_channels=3, out_channels=1).to(DEVICE)\nmodel.load_state_dict(checkpoint['state_dict'])\n\nmodel.eval()\n\n# Predictions\nall_predictions = []\nfor img_names, x in tqdm(test_loader):\n    x = x.to(DEVICE)\n    with torch.no_grad():\n        preds = torch.sigmoid(model(x))\n        preds = (preds > THRESHOLD).float()   \n    preds = TF.resize(\n        preds, size=(1280, 1918), interpolation=TF.InterpolationMode.NEAREST\n    )\n    \n    # Encoding\n    for idx in range(len(img_names)):\n        encoding = rle_encode(preds[idx].squeeze().cpu())\n        all_predictions.append([img_names[idx], encoding])","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree(WORKING_DIR + 'test')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub = pd.DataFrame(all_predictions)\nsub.columns = ['img', 'rle_mask']\nsub.to_csv(os.path.join(WORKING_DIR, 'submission.csv'), index=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom IPython.display import FileLink\n\nos.chdir('/kaggle/working')\nFileLink('submission.csv')","metadata":{},"execution_count":null,"outputs":[]}]}