{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","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":30627,"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":{"execution":{"iopub.status.busy":"2023-12-16T21:35:25.612857Z","iopub.execute_input":"2023-12-16T21:35:25.613492Z","iopub.status.idle":"2023-12-16T21:35:25.619977Z","shell.execute_reply.started":"2023-12-16T21:35:25.613438Z","shell.execute_reply":"2023-12-16T21:35:25.619064Z"},"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":"2023-12-16T21:35:25.624239Z","iopub.execute_input":"2023-12-16T21:35:25.624512Z","iopub.status.idle":"2023-12-16T21:35:25.634390Z","shell.execute_reply.started":"2023-12-16T21:35:25.624480Z","shell.execute_reply":"2023-12-16T21:35:25.633612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sorted(os.listdir(WORKING_DIR + 'train'))[:10]","metadata":{"execution":{"iopub.status.busy":"2023-12-16T21:35:25.638652Z","iopub.execute_input":"2023-12-16T21:35:25.638946Z","iopub.status.idle":"2023-12-16T21:35:25.651299Z","shell.execute_reply.started":"2023-12-16T21:35:25.638897Z","shell.execute_reply":"2023-12-16T21:35:25.650495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sorted(os.listdir(WORKING_DIR + 'train_masks'))[:10]","metadata":{"execution":{"iopub.status.busy":"2023-12-16T21:35:25.653042Z","iopub.execute_input":"2023-12-16T21:35:25.653686Z","iopub.status.idle":"2023-12-16T21:35:25.664873Z","shell.execute_reply.started":"2023-12-16T21:35:25.653653Z","shell.execute_reply":"2023-12-16T21:35:25.663920Z"},"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-12-16T21:35:25.665973Z","iopub.execute_input":"2023-12-16T21:35:25.666302Z","iopub.status.idle":"2023-12-16T21:35:25.675364Z","shell.execute_reply.started":"2023-12-16T21:35:25.666270Z","shell.execute_reply":"2023-12-16T21:35:25.674515Z"},"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)\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.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)\n    \ndef 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\n\nif __name__ == \"__main__\":\n    test()","metadata":{"execution":{"iopub.status.busy":"2023-12-16T21:35:25.678883Z","iopub.execute_input":"2023-12-16T21:35:25.679130Z","iopub.status.idle":"2023-12-16T21:35:26.757113Z","shell.execute_reply.started":"2023-12-16T21:35:25.679108Z","shell.execute_reply":"2023-12-16T21:35:26.756162Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torchvision\n# from dataset import CarvanaDataset\nfrom torch.utils.data import DataLoader\n\ndef 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,\n    train_maskdir,\n    val_dir,\n    val_maskdir,\n    batch_size,\n    train_transform,\n    val_transform,\n    num_workers=4,\n    pin_memory=True,\n):\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":"2023-12-16T21:35:26.759376Z","iopub.execute_input":"2023-12-16T21:35:26.759801Z","iopub.status.idle":"2023-12-16T21:35:26.773721Z","shell.execute_reply.started":"2023-12-16T21:35:26.759766Z","shell.execute_reply":"2023-12-16T21:35:26.772818Z"},"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# from model import UNET\n# from utils import (\n#     load_checkpoint,\n#     save_checkpoint,\n#     get_loaders,\n#     check_accuracy,\n#     save_predictions_as_imgs,\n# )\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/'\n\ndef 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())\n\ndef main():\n    train_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\n    val_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\n    model = UNET(in_channels=3, out_channels=1).to(DEVICE)\n    loss_fn = nn.BCEWithLogitsLoss()\n    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n\n    train_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\n    if LOAD_MODEL:\n        load_checkpoint(torch.load('my_checkpoint.pth.tar'), model)\n        check_accuracy(val_loader, model, device=DEVICE)\n\n    scaler = torch.cuda.amp.GradScaler()\n\n    for 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        )\n\nif __name__ == '__main__':\n    main()","metadata":{"execution":{"iopub.status.busy":"2023-12-16T21:35:26.774952Z","iopub.execute_input":"2023-12-16T21:35:26.775256Z","iopub.status.idle":"2023-12-16T21:35:37.681413Z","shell.execute_reply.started":"2023-12-16T21:35:26.775222Z","shell.execute_reply":"2023-12-16T21:35:37.679507Z"},"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)","metadata":{"execution":{"iopub.status.busy":"2023-12-17T09:17:48.320997Z","iopub.execute_input":"2023-12-17T09:17:48.321802Z","iopub.status.idle":"2023-12-17T09:17:48.328751Z","shell.execute_reply.started":"2023-12-17T09:17:48.321737Z","shell.execute_reply":"2023-12-17T09:17:48.327727Z"},"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":"2023-12-16T22:43:46.199716Z","iopub.execute_input":"2023-12-16T22:43:46.200111Z","iopub.status.idle":"2023-12-16T22:46:45.518099Z","shell.execute_reply.started":"2023-12-16T22:43:46.200082Z","shell.execute_reply":"2023-12-16T22:46:45.517210Z"},"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":{"iopub.status.busy":"2023-12-17T09:17:56.106190Z","iopub.execute_input":"2023-12-17T09:17:56.106534Z","iopub.status.idle":"2023-12-17T10:24:10.424943Z","shell.execute_reply.started":"2023-12-17T09:17:56.106507Z","shell.execute_reply":"2023-12-17T10:24:10.423988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree(WORKING_DIR + 'test')","metadata":{"execution":{"iopub.status.busy":"2023-12-17T10:41:38.198496Z","iopub.execute_input":"2023-12-17T10:41:38.198953Z","iopub.status.idle":"2023-12-17T10:41:43.175090Z","shell.execute_reply.started":"2023-12-17T10:41:38.198918Z","shell.execute_reply":"2023-12-17T10:41:43.174057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Submission\nsub = pd.DataFrame(all_predictions)\nsub.columns = ['img', 'rle_mask']\nsub.to_csv(os.path.join(WORKING_DIR, 'submission.csv'), index=False)","metadata":{"execution":{"iopub.status.busy":"2023-12-17T10:41:43.177039Z","iopub.execute_input":"2023-12-17T10:41:43.177469Z","iopub.status.idle":"2023-12-17T10:42:18.767710Z","shell.execute_reply.started":"2023-12-17T10:41:43.177441Z","shell.execute_reply":"2023-12-17T10:42:18.766666Z"},"trusted":true},"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":{"iopub.status.busy":"2023-12-17T10:42:18.768966Z","iopub.execute_input":"2023-12-17T10:42:18.769263Z","iopub.status.idle":"2023-12-17T10:42:18.776820Z","shell.execute_reply.started":"2023-12-17T10:42:18.769236Z","shell.execute_reply":"2023-12-17T10:42:18.775834Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}