vesuvius-challenge-example-submission-e322cd.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["### Imports"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-04-02T14:30:09.884568Z","iopub.status.busy":"2023-04-02T14:30:09.883896Z","iopub.status.idle":"2023-04-02T14:30:09.894539Z","shell.execute_reply":"2023-04-02T14:30:09.893150Z","shell.execute_reply.started":"2023-04-02T14:30:09.884530Z"},"trusted":true},"outputs":[],"source":["import os\n","import gc\n","import glob\n","import json\n","from collections import defaultdict\n","import multiprocessing as mp\n","from pathlib import Path\n","from types import SimpleNamespace\n","from typing import Dict, List, Optional, Tuple\n","import warnings\n","\n","import matplotlib.pyplot as plt\n","import matplotlib.patches as patches\n","import numpy as np\n","import pandas as pd\n","import PIL.Image as Image\n","from sklearn.metrics import fbeta_score\n","from sklearn.exceptions import UndefinedMetricWarning\n","import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","import torch.utils.data as thd\n","from tqdm import tqdm"]},{"cell_type":"markdown","metadata":{},"source":["### Set up data"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-04-02T14:30:09.897537Z","iopub.status.busy":"2023-04-02T14:30:09.897059Z","iopub.status.idle":"2023-04-02T14:30:09.924316Z","shell.execute_reply":"2023-04-02T14:30:09.923315Z","shell.execute_reply.started":"2023-04-02T14:30:09.897500Z"},"trusted":true},"outputs":[],"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()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-04-02T14:30:09.928018Z","iopub.status.busy":"2023-04-02T14:30:09.927433Z","iopub.status.idle":"2023-04-02T14:30:09.949383Z","shell.execute_reply":"2023-04-02T14:30:09.948375Z","shell.execute_reply.started":"2023-04-02T14:30:09.927991Z"},"trusted":true},"outputs":[],"source":["base_path = Path(\"/workspace/\")\n","train_path = base_path / \"train\"\n","all_fragments = sorted([f.name for f in train_path.iterdir()])\n","print(\"All fragments:\", all_fragments)\n","# Due to limited memory on Kaggle, we can only load 1 full fragment\n","train_fragments = [train_path / fragment_name for fragment_name in [\"1\"]]\n","train_fragments"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-04-02T14:30:09.952646Z","iopub.status.busy":"2023-04-02T14:30:09.951765Z"},"trusted":true},"outputs":[],"source":["train_dset = SubvolumeDataset(fragments=train_fragments, voxel_shape=(48, 64, 64), filter_edge_pixels=True)\n","print(\"Num items (pixels)\", len(train_dset))"]},{"cell_type":"markdown","metadata":{},"source":["#### Sanity check "]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["index = 6136130\n","train_dset.plot_label(index, figsize=(16, 10))"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["BATCH_SIZE = 32\n","train_loader = thd.DataLoader(train_dset, batch_size=BATCH_SIZE, shuffle=True)\n","print(\"Num batches:\", len(train_loader))"]},{"cell_type":"markdown","metadata":{},"source":["### Set up model"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"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)"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["model = InkDetector().to(DEVICE)"]},{"cell_type":"markdown","metadata":{},"source":["### Train"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["TRAINING_STEPS = 60000\n","LEARNING_RATE = 1e-3\n","TRAIN_RUN = True # To avoid re-running when saving the notebook"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["warnings.simplefilter('ignore', UndefinedMetricWarning)"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"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(), \"/workspace/model.pt\")\n","\n","else:\n","    model_weights = torch.load(\"/workspace/model.pt\")\n","    model.load_state_dict(model_weights)"]},{"cell_type":"markdown","metadata":{},"source":["### Evaluate"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["# Clear memory before loading test fragments\n","train_dset.labels = None\n","train_dset.image_stacks = []\n","del train_loader, train_dset\n","gc.collect()"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["test_path = base_path / \"test\"\n","test_fragments = [train_path / fragment_name for fragment_name in test_path.iterdir()]\n","print(\"All fragments:\", test_fragments)"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["pred_images = []\n","model.eval()\n","for 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)"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["plt.imshow(pred_images[1], cmap='gray')"]},{"cell_type":"markdown","metadata":{},"source":["### Submission"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"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), ())))"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["submission = defaultdict(list)\n","for 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","\n","pd.DataFrame.from_dict(submission).to_csv(\"/workspace/submission.csv\", index=False)"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["pd.DataFrame.from_dict(submission)"]}],"metadata":{"kernelspec":{"display_name":"venv","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.13"},"vscode":{"interpreter":{"hash":"1579f81f26a28ddd3095484c461e47df81abfb941bbf90f730815a9b45962d81"}}},"nbformat":4,"nbformat_minor":4}