{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":47317,"databundleVersionId":5799376,"sourceType":"competition"}],"dockerImageVersionId":30408,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nimport glob\nimport json\nfrom collections import defaultdict\nimport multiprocessing as mp\nfrom pathlib import Path\nfrom types import SimpleNamespace\nfrom typing import Dict, List, Optional, Tuple\nimport warnings\n\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nimport numpy as np\nimport pandas as pd\nimport PIL.Image as Image\nfrom sklearn.metrics import fbeta_score\nfrom sklearn.exceptions import UndefinedMetricWarning\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.utils.data as thd\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:50:08.952828Z","iopub.execute_input":"2023-12-01T13:50:08.953135Z","iopub.status.idle":"2023-12-01T13:50:12.208710Z","shell.execute_reply.started":"2023-12-01T13:50:08.953105Z","shell.execute_reply":"2023-12-01T13:50:12.207491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Set up data","metadata":{}},{"cell_type":"code","source":"class SubvolumeDataset(thd.Dataset):\n    def __init__(\n        self,\n        fragments: List[Path],\n        voxel_shape: Tuple[int, int, int],\n        load_inklabels: bool = True,\n        filter_edge_pixels: bool = False,\n    ):\n        self.fragments = sorted(map(lambda path: path.resolve(), fragments))\n        self.voxel_shape = voxel_shape\n        self.load_inklabels = load_inklabels\n        self.filter_edge_pixels = filter_edge_pixels\n\n        # Load sequentially\n        labels = []\n        image_stacks = []\n        valid_pixels = []\n        for fragment_id, fragment_path in enumerate(self.fragments):\n            fragment_path = fragment_path.resolve()  # absolute path\n            mask = np.array(Image.open(str(fragment_path / \"mask.png\")).convert(\"1\"))\n\n            surface_volume_paths = sorted(\n                (fragment_path / \"surface_volume\").rglob(\"*.tif\")\n            )\n            z_dim, y_dim, x_dim = voxel_shape\n\n            z_mid = len(surface_volume_paths) // 2\n            z_start, z_end = z_mid - z_dim // 2, z_mid + z_dim // 2\n\n            # we don't convert to torch since it doesn't support uint16\n            images = [\n                np.array(Image.open(fn)) for fn in surface_volume_paths[z_start:z_end]\n            ]\n            image_stack = np.stack(images, axis=0)\n            image_stacks.append(image_stack)\n\n            pixels = np.stack(np.where(mask == 1), axis=1).astype(np.uint16)\n            if filter_edge_pixels:\n                height, width = mask.shape\n                mask_y = np.logical_or(\n                    pixels[:, 0] < y_dim // 2, pixels[:, 0] >= height - y_dim // 2\n                )\n                mask_x = np.logical_or(\n                    pixels[:, 1] < x_dim // 2, pixels[:, 1] >= width - x_dim // 2\n                )\n                pixel_mask = np.logical_or(mask_y, mask_x)\n                pixels = pixels[~pixel_mask]\n            # encode fragment ID\n            fragment_ids = np.full_like(pixels[:, 0:1], fragment_id)\n            pixels = np.concatenate((pixels, fragment_ids), axis=1)\n            valid_pixels.append(pixels)\n\n            if load_inklabels:\n                # binary mask can be stored as np.bool\n                inklabels = (\n                    np.array(Image.open(str(fragment_path / \"inklabels.png\"))) > 0\n                )\n                labels.append(inklabels)\n\n            print(f\"Loaded fragment {fragment_path} on {os.getpid()}\")\n\n        self.labels = labels\n        self.image_stacks = image_stacks\n        self.pixels = np.concatenate(valid_pixels).reshape(\n            -1, valid_pixels[0].shape[-1]\n        )\n\n    def __len__(self):\n        return len(self.pixels)\n\n    def __getitem__(self, index):\n        center_y, center_x, fragment_id = self.pixels[index]\n        z_dim, y_dim, x_dim = self.voxel_shape\n        image_stack = self.image_stacks[fragment_id]\n        _, height, width = image_stack.shape\n\n        # pad with zeros if necessary\n        if (\n            center_y < y_dim // 2\n            or center_x < x_dim // 2\n            or center_y + y_dim // 2 >= height\n            or center_x + x_dim // 2 >= width\n        ):\n            # calculate the upper-left corner of the sub-volume\n            y_start = max(center_y - y_dim // 2, 0)\n            x_start = max(center_x - x_dim // 2, 0)\n\n            # calculate the lower-right corner of the sub-volume\n            y_end = min(center_y + y_dim // 2, height)\n            x_end = min(center_x + x_dim // 2, width)\n\n            subvolume = np.zeros(self.voxel_shape, dtype=np.float32)\n\n            pad_y_start = max(y_dim // 2 - center_y, 0)\n            pad_x_start = max(x_dim // 2 - center_x, 0)\n\n            pad_y_end = min(height + y_dim // 2 - center_y, y_dim)\n            pad_x_end = min(width + x_dim // 2 - center_x, x_dim)\n\n            subvolume[:, pad_y_start:pad_y_end, pad_x_start:pad_x_end] = (\n                image_stack[:, y_start:y_end, x_start:x_end].astype(np.float32) / 65535\n            )\n\n        else:\n            subvolume = (\n                image_stack[\n                    :,\n                    center_y - y_dim // 2 : center_y + y_dim // 2,\n                    center_x - x_dim // 2 : center_x + x_dim // 2,\n                ]\n            ).astype(np.float32) / 65535\n        if self.load_inklabels:\n            inklabel = float(self.labels[fragment_id][center_y, center_x])\n        else:\n            inklabel = -1.0\n\n        return torch.from_numpy(subvolume).unsqueeze(0), torch.FloatTensor([inklabel])\n\n    def plot_label(self, index, **kwargs):\n        pixel = self.pixels[index]\n        label = self.labels[pixel[-1]]\n\n        print(\"Index:\", index)\n        print(\"Pixel:\", pixel)\n        print(\"Label:\", int(label[pixel[0], pixel[1]]))\n\n        if isinstance(label, torch.Tensor):\n            label = label.numpy()\n\n        fig, ax = plt.subplots(**kwargs)\n        ax.imshow(label, cmap=\"gray\")\n\n        y, x, _ = pixel\n        _, y_dim, x_dim = self.voxel_shape\n        x_min = x - (x_dim // 2)\n        x_max = x + (x_dim // 2)\n        y_min = y - (y_dim // 2)\n        y_max = y + (y_dim // 2)\n\n        rect = plt.Rectangle(\n            (x_min, y_min), x_dim, y_dim, linewidth=2, edgecolor=\"y\", facecolor=\"none\"\n        )\n        ax.add_patch(rect)\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:50:12.212054Z","iopub.execute_input":"2023-12-01T13:50:12.212722Z","iopub.status.idle":"2023-12-01T13:50:12.249649Z","shell.execute_reply.started":"2023-12-01T13:50:12.212668Z","shell.execute_reply":"2023-12-01T13:50:12.248207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_path = Path(\"/kaggle/input/vesuvius-challenge/\")\ntrain_path = base_path / \"train\"\nall_fragments = sorted([f.name for f in train_path.iterdir()])\nprint(\"All fragments:\", all_fragments)\n# Due to limited memory on Kaggle, we can only load 1 full fragment\ntrain_fragments = [train_path / fragment_name for fragment_name in [\"1\"]]\ntrain_fragments","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:50:12.251235Z","iopub.execute_input":"2023-12-01T13:50:12.251585Z","iopub.status.idle":"2023-12-01T13:50:12.282125Z","shell.execute_reply.started":"2023-12-01T13:50:12.251553Z","shell.execute_reply":"2023-12-01T13:50:12.281073Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ntrain_dset = SubvolumeDataset(fragments=train_fragments, voxel_shape=(48, 64, 64), filter_edge_pixels=True)\nprint(\"Num items (pixels)\", len(train_dset))","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:50:12.285463Z","iopub.execute_input":"2023-12-01T13:50:12.285861Z","iopub.status.idle":"2023-12-01T13:51:42.001814Z","shell.execute_reply.started":"2023-12-01T13:50:12.285828Z","shell.execute_reply":"2023-12-01T13:51:42.000617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Sanity check ","metadata":{}},{"cell_type":"code","source":"index = 6136130\ntrain_dset.plot_label(index, figsize=(16, 10))","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:42.003088Z","iopub.execute_input":"2023-12-01T13:51:42.003402Z","iopub.status.idle":"2023-12-01T13:51:45.017514Z","shell.execute_reply.started":"2023-12-01T13:51:42.003371Z","shell.execute_reply":"2023-12-01T13:51:45.016337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 32\ntrain_loader = thd.DataLoader(train_dset, batch_size=BATCH_SIZE, shuffle=True)\nprint(\"Num batches:\", len(train_loader))","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:45.018998Z","iopub.execute_input":"2023-12-01T13:51:45.019402Z","iopub.status.idle":"2023-12-01T13:51:45.026456Z","shell.execute_reply.started":"2023-12-01T13:51:45.019362Z","shell.execute_reply":"2023-12-01T13:51:45.025322Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Set up model","metadata":{}},{"cell_type":"code","source":"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:45.027982Z","iopub.execute_input":"2023-12-01T13:51:45.028519Z","iopub.status.idle":"2023-12-01T13:51:45.124538Z","shell.execute_reply.started":"2023-12-01T13:51:45.028476Z","shell.execute_reply":"2023-12-01T13:51:45.123304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class InkDetector(torch.nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        filters = [16, 32, 64]\n        paddings = [1, 1, 1]\n        kernel_sizes = [3, 3, 3]\n        strides = [2, 2, 2]\n        \n        layers = []\n        in_channels = 1\n        for num_filters, padding, kernel_size, stride in zip(filters, paddings, kernel_sizes, strides):\n            layers.extend([\n                nn.Conv3d(\n                    in_channels=in_channels,\n                    out_channels=num_filters,\n                    kernel_size=kernel_size,\n                    stride=stride,\n                    padding=padding,\n                ),\n                nn.ReLU(inplace=True),\n                torch.nn.BatchNorm3d(num_features=num_filters)\n            ])\n            in_channels = num_filters\n        layers.append(nn.AdaptiveAvgPool3d(1))\n        layers.append(nn.Flatten())\n\n        self.encoder = nn.Sequential(*layers)\n        self.decoder = nn.Sequential(\n            nn.Linear(in_channels, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        features = self.encoder(x)\n        return self.decoder(features)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:45.126040Z","iopub.execute_input":"2023-12-01T13:51:45.126825Z","iopub.status.idle":"2023-12-01T13:51:45.138819Z","shell.execute_reply.started":"2023-12-01T13:51:45.126768Z","shell.execute_reply":"2023-12-01T13:51:45.137763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = InkDetector().to(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:45.140208Z","iopub.execute_input":"2023-12-01T13:51:45.140558Z","iopub.status.idle":"2023-12-01T13:51:47.656390Z","shell.execute_reply.started":"2023-12-01T13:51:45.140527Z","shell.execute_reply":"2023-12-01T13:51:47.655441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Train","metadata":{}},{"cell_type":"code","source":"TRAINING_STEPS = 60000\nLEARNING_RATE = 1e-3\nTRAIN_RUN = True # To avoid re-running when saving the notebook","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:47.660789Z","iopub.execute_input":"2023-12-01T13:51:47.661098Z","iopub.status.idle":"2023-12-01T13:51:47.665821Z","shell.execute_reply.started":"2023-12-01T13:51:47.661069Z","shell.execute_reply":"2023-12-01T13:51:47.664717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"warnings.simplefilter('ignore', UndefinedMetricWarning)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:47.667134Z","iopub.execute_input":"2023-12-01T13:51:47.667448Z","iopub.status.idle":"2023-12-01T13:51:47.674400Z","shell.execute_reply.started":"2023-12-01T13:51:47.667418Z","shell.execute_reply":"2023-12-01T13:51:47.673243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if TRAIN_RUN:\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE)\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\n    model.train()\n    running_loss = 0.0\n    running_accuracy = 0.0\n    running_fbeta = 0.0\n    denom = 0\n    pbar = tqdm(enumerate(train_loader), total=TRAINING_STEPS)\n    for i, (subvolumes, inklabels) in pbar:\n        if i >= TRAINING_STEPS:\n            break\n        optimizer.zero_grad()\n        outputs = model(subvolumes.to(DEVICE))\n        loss = criterion(outputs, inklabels.to(DEVICE))\n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n        pred_ink = outputs.detach().sigmoid().gt(0.4).cpu().int()\n        accuracy = (pred_ink == inklabels).sum().float().div(inklabels.size(0))\n        running_fbeta += fbeta_score(inklabels.view(-1).numpy(), pred_ink.view(-1).numpy(), beta=0.5)\n        running_accuracy += accuracy.item()\n        running_loss += loss.item()\n        denom += 1\n        pbar.set_postfix({\"Loss\": running_loss / denom, \"Accuracy\": running_accuracy / denom, \"Fbeta@0.5\": running_fbeta / denom})\n        if (i + 1) % 500 == 0:\n            running_loss = 0.\n            running_accuracy = 0.\n            running_fbeta = 0.\n            denom = 0\n\n    torch.save(model.state_dict(), \"/kaggle/working/model.pt\")\n\nelse:\n    model_weights = torch.load(\"/kaggle/working/model.pt\")\n    model.load_state_dict(model_weights)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T13:51:47.675826Z","iopub.execute_input":"2023-12-01T13:51:47.676226Z","iopub.status.idle":"2023-12-01T15:29:05.863078Z","shell.execute_reply.started":"2023-12-01T13:51:47.676161Z","shell.execute_reply":"2023-12-01T15:29:05.861885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Evaluate","metadata":{}},{"cell_type":"code","source":"# Clear memory before loading test fragments\ntrain_dset.labels = None\ntrain_dset.image_stacks = []\ndel train_loader, train_dset\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-12-01T15:29:05.865580Z","iopub.execute_input":"2023-12-01T15:29:05.866192Z","iopub.status.idle":"2023-12-01T15:29:06.273561Z","shell.execute_reply.started":"2023-12-01T15:29:05.866156Z","shell.execute_reply":"2023-12-01T15:29:06.271958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_path = base_path / \"test\"\ntest_fragments = [test_path / fragment_name for fragment_name in test_path.iterdir()]\nprint(\"All fragments:\", test_fragments)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T15:29:06.275026Z","iopub.execute_input":"2023-12-01T15:29:06.275550Z","iopub.status.idle":"2023-12-01T15:29:06.291030Z","shell.execute_reply.started":"2023-12-01T15:29:06.275505Z","shell.execute_reply":"2023-12-01T15:29:06.290093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_images = []\nmodel.eval()\nfor test_fragment in test_fragments:\n    outputs = []\n    eval_dset = SubvolumeDataset(fragments=[test_fragment], voxel_shape=(48, 64, 64), load_inklabels=False)\n    eval_loader = thd.DataLoader(eval_dset, batch_size=BATCH_SIZE, shuffle=False)\n    with torch.no_grad():\n        for i, (subvolumes, _) in enumerate(tqdm(eval_loader)):\n            output = model(subvolumes.to(DEVICE)).view(-1).sigmoid().cpu().numpy()\n            outputs.append(output)\n    # we only load 1 fragment at a time\n    image_shape = eval_dset.image_stacks[0].shape[1:]\n    eval_dset.labels = None\n    eval_dset.image_stacks = None\n    del eval_loader\n    gc.collect()\n\n    pred_image = np.zeros(image_shape, dtype=np.uint8)\n    outputs = np.concatenate(outputs)\n    for (y, x, _), prob in zip(eval_dset.pixels[:outputs.shape[0]], outputs):\n        pred_image[y ,x] = prob > 0.4\n    pred_images.append(pred_image)\n    \n    eval_dset.pixels = None\n    del eval_dset\n    gc.collect()\n    print(\"Finished\", test_fragment)","metadata":{"execution":{"iopub.status.busy":"2023-12-01T15:29:06.292356Z","iopub.execute_input":"2023-12-01T15:29:06.293192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(pred_images[1], cmap='gray')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Submission","metadata":{}},{"cell_type":"code","source":"def rle(output):\n    flat_img = np.where(output > 0.4, 1, 0).astype(np.uint8)\n    starts = np.array((flat_img[:-1] == 0) & (flat_img[1:] == 1))\n    ends = np.array((flat_img[:-1] == 1) & (flat_img[1:] == 0))\n    starts_ix = np.where(starts)[0] + 2\n    ends_ix = np.where(ends)[0] + 2\n    lengths = ends_ix - starts_ix\n    return \" \".join(map(str, sum(zip(starts_ix, lengths), ())))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = defaultdict(list)\nfor fragment_id, fragment_name in enumerate(test_fragments):\n    submission[\"Id\"].append(fragment_name.name)\n    submission[\"Predicted\"].append(rle(pred_images[fragment_id]))\n\npd.DataFrame.from_dict(submission).to_csv(\"/kaggle/working/submission.csv\", index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.DataFrame.from_dict(submission)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}