{"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":"markdown","source":"**Note:** Using [Vesuvius Challenge: Example submission](https://www.kaggle.com/code/danielhavir/vesuvius-challenge-example-submission) as source, I modified the reading function to use tifffile optimal readers to speed up the dataloader definition. In the *Data import Optimization* section there is a brief explanation of how the speed up is done.","metadata":{}},{"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 tifffile\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-03-31T08:12:46.635343Z","iopub.execute_input":"2023-03-31T08:12:46.635634Z","iopub.status.idle":"2023-03-31T08:12:50.647045Z","shell.execute_reply.started":"2023-03-31T08:12:46.635607Z","shell.execute_reply":"2023-03-31T08:12:50.645937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data import Optimization\n\nThe library *tifffile* can be used to optimize the file reading to seconds.\n\nThere is a class called *tifffile.TiffSequence* which allows you to optimize the data reading using memory references. This class can take a filelist and generate a numpy array without reading the images, all done by references. This can be done using *asarray(out='memmap')* method to the image sequence object.\n\nIn the next example, this class is used to read the reference of the first 3 tiff files from the train image 1. This code shows how the object has three layers, but it doesn't load the data until needed. So, its possible to generate the datareader with all three training images.","metadata":{}},{"cell_type":"code","source":"%%time\npath = '/kaggle/input/vesuvius-challenge/train/1/surface_volume/'\nall_files = np.sort(glob.glob(path + '*'))\n\ntiffsquence = tifffile.TiffSequence(\n    files=all_files[:3]\n)\nimg_pointer = tiffsquence.asarray(out='memmap')\nprint('Squence shape:', img_pointer.shape)","metadata":{"execution":{"iopub.status.busy":"2023-03-31T08:12:50.654905Z","iopub.execute_input":"2023-03-31T08:12:50.655492Z","iopub.status.idle":"2023-03-31T08:12:53.666978Z","shell.execute_reply.started":"2023-03-31T08:12:50.655452Z","shell.execute_reply":"2023-03-31T08:12:53.665739Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nplt.imshow(img_pointer[1, 2000:4000, 2000:4000], cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-31T08:13:09.433871Z","iopub.execute_input":"2023-03-31T08:13:09.434975Z","iopub.status.idle":"2023-03-31T08:13:09.958597Z","shell.execute_reply.started":"2023-03-31T08:13:09.434930Z","shell.execute_reply":"2023-03-31T08:13:09.957596Z"},"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_pointers_dict = {}\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\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 generate the references for the image\n            tiffsquence = tifffile.TiffSequence(\n                files=surface_volume_paths[z_start:z_end]\n            )\n            image_pointers_dict[fragment_id] = tiffsquence.asarray(out='memmap')\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_pointers_dict = image_pointers_dict\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_pointer = self.image_pointers_dict[fragment_id]\n        _, height, width = image_pointer.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_pointer[:, y_start:y_end, x_start:x_end].astype(np.float32) / 65535\n            )\n\n        else:\n            subvolume = (\n                image_pointer[\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-03-31T06:59:46.362997Z","iopub.execute_input":"2023-03-31T06:59:46.363403Z","iopub.status.idle":"2023-03-31T06:59:46.403457Z","shell.execute_reply.started":"2023-03-31T06:59:46.363360Z","shell.execute_reply":"2023-03-31T06:59:46.402401Z"},"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 all_fragments]  # [\"1\"]]\ntrain_fragments","metadata":{"execution":{"iopub.status.busy":"2023-03-31T06:59:46.407916Z","iopub.execute_input":"2023-03-31T06:59:46.410302Z","iopub.status.idle":"2023-03-31T06:59:46.427583Z","shell.execute_reply.started":"2023-03-31T06:59:46.410260Z","shell.execute_reply":"2023-03-31T06:59:46.426609Z"},"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-03-31T06:59:46.431700Z","iopub.execute_input":"2023-03-31T06:59:46.434121Z","iopub.status.idle":"2023-03-31T07:03:37.075929Z","shell.execute_reply.started":"2023-03-31T06:59:46.434075Z","shell.execute_reply":"2023-03-31T07:03:37.074671Z"},"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-03-31T07:03:37.077736Z","iopub.execute_input":"2023-03-31T07:03:37.078166Z","iopub.status.idle":"2023-03-31T07:03:40.275062Z","shell.execute_reply.started":"2023-03-31T07:03:37.078126Z","shell.execute_reply":"2023-03-31T07:03:40.274028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 8\ntrain_loader = thd.DataLoader(train_dset, batch_size=BATCH_SIZE, shuffle=True)\nprint(\"Num batches:\", len(train_loader))","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:03:40.276068Z","iopub.execute_input":"2023-03-31T07:03:40.277308Z","iopub.status.idle":"2023-03-31T07:03:40.284379Z","shell.execute_reply.started":"2023-03-31T07:03:40.277272Z","shell.execute_reply":"2023-03-31T07:03:40.282225Z"},"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-03-31T07:03:40.286096Z","iopub.execute_input":"2023-03-31T07:03:40.286688Z","iopub.status.idle":"2023-03-31T07:03:40.383190Z","shell.execute_reply.started":"2023-03-31T07:03:40.286653Z","shell.execute_reply":"2023-03-31T07:03:40.382078Z"},"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-03-31T07:03:40.387723Z","iopub.execute_input":"2023-03-31T07:03:40.388041Z","iopub.status.idle":"2023-03-31T07:03:40.398446Z","shell.execute_reply.started":"2023-03-31T07:03:40.388013Z","shell.execute_reply":"2023-03-31T07:03:40.397131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = InkDetector().to(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:03:40.400282Z","iopub.execute_input":"2023-03-31T07:03:40.401098Z","iopub.status.idle":"2023-03-31T07:03:43.364678Z","shell.execute_reply.started":"2023-03-31T07:03:40.401061Z","shell.execute_reply":"2023-03-31T07:03:43.363669Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Train","metadata":{}},{"cell_type":"code","source":"TRAINING_STEPS = 5000\nLEARNING_RATE = 1e-3\nTRAIN_RUN = True # To avoid re-running when saving the notebook","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:03:43.367293Z","iopub.execute_input":"2023-03-31T07:03:43.368352Z","iopub.status.idle":"2023-03-31T07:03:43.374913Z","shell.execute_reply.started":"2023-03-31T07:03:43.368283Z","shell.execute_reply":"2023-03-31T07:03:43.373803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"warnings.simplefilter('ignore', UndefinedMetricWarning)","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:03:43.376419Z","iopub.execute_input":"2023-03-31T07:03:43.377238Z","iopub.status.idle":"2023-03-31T07:03:44.720480Z","shell.execute_reply.started":"2023-03-31T07:03:43.377189Z","shell.execute_reply":"2023-03-31T07:03:44.719376Z"},"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) % 100 == 0:\n            running_loss = 0.\n            running_accuracy = 0.\n            running_fbeta = 0.\n            denom = 0\n            torch.save(model.state_dict(), \"/kaggle/working/model_checkpoint.pt\")\n\n    torch.save(model.state_dict(), \"/kaggle/working/model.pt\")\n\nelse:\n    model_weights = torch.load(\"/kaggle/working/model_checkpoint.pt\")\n    model.load_state_dict(model_weights)","metadata":{"_kg_hide-output":false,"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_pointers_dict = {}\ndel train_loader, train_dset\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:56:19.591835Z","iopub.execute_input":"2023-03-31T07:56:19.592797Z","iopub.status.idle":"2023-03-31T07:56:21.930461Z","shell.execute_reply.started":"2023-03-31T07:56:19.592756Z","shell.execute_reply":"2023-03-31T07:56:21.927381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_path = base_path / \"test\"\ntest_fragments = [train_path / fragment_name for fragment_name in test_path.iterdir()]\nprint(\"All fragments:\", test_fragments)","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:56:43.367474Z","iopub.execute_input":"2023-03-31T07:56:43.368070Z","iopub.status.idle":"2023-03-31T07:56:43.376280Z","shell.execute_reply.started":"2023-03-31T07:56:43.368032Z","shell.execute_reply":"2023-03-31T07:56:43.375000Z"},"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_pointers_dict[0].shape[1:]\n    eval_dset.labels = None\n    eval_dset.image_pointers_dict = {}\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":{"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":{"execution":{"iopub.status.busy":"2023-03-31T07:55:48.674555Z","iopub.status.idle":"2023-03-31T07:55:48.675424Z","shell.execute_reply.started":"2023-03-31T07:55:48.675158Z","shell.execute_reply":"2023-03-31T07:55:48.675186Z"},"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":{"execution":{"iopub.status.busy":"2023-03-31T07:55:48.676952Z","iopub.status.idle":"2023-03-31T07:55:48.682780Z","shell.execute_reply.started":"2023-03-31T07:55:48.682499Z","shell.execute_reply":"2023-03-31T07:55:48.682528Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.DataFrame.from_dict(submission)","metadata":{"execution":{"iopub.status.busy":"2023-03-31T07:55:48.684193Z","iopub.status.idle":"2023-03-31T07:55:48.684973Z","shell.execute_reply.started":"2023-03-31T07:55:48.684712Z","shell.execute_reply":"2023-03-31T07:55:48.684737Z"},"trusted":true},"execution_count":null,"outputs":[]}]}