{"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":"code","source":"from pathlib import Path\nimport numpy as np\nfrom PIL import Image\nfrom collections import defaultdict\nimport pandas as pd\n\nimport 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","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:31:25.722523Z","iopub.execute_input":"2023-05-12T07:31:25.723683Z","iopub.status.idle":"2023-05-12T07:31:25.731907Z","shell.execute_reply.started":"2023-05-12T07:31:25.723644Z","shell.execute_reply":"2023-05-12T07:31:25.730183Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import 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-05-12T07:35:27.672828Z","iopub.execute_input":"2023-05-12T07:35:27.673189Z","iopub.status.idle":"2023-05-12T07:35:27.679611Z","shell.execute_reply.started":"2023-05-12T07:35:27.673159Z","shell.execute_reply":"2023-05-12T07:35:27.678414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_path = Path(\"/kaggle/input/vesuvius-challenge-ink-detection/\")\ntrain_path = base_path / \"train\"\ntest_path = base_path / \"test\"","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:31:32.102669Z","iopub.execute_input":"2023-05-12T07:31:32.103022Z","iopub.status.idle":"2023-05-12T07:31:32.107486Z","shell.execute_reply.started":"2023-05-12T07:31:32.102993Z","shell.execute_reply":"2023-05-12T07:31:32.106564Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 32","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:35:34.028943Z","iopub.execute_input":"2023-05-12T07:35:34.029930Z","iopub.status.idle":"2023-05-12T07:35:34.039211Z","shell.execute_reply.started":"2023-05-12T07:35:34.029887Z","shell.execute_reply":"2023-05-12T07:35:34.038150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-05-12T07:31:35.178558Z","iopub.execute_input":"2023-05-12T07:31:35.178911Z","iopub.status.idle":"2023-05-12T07:31:35.574094Z","shell.execute_reply.started":"2023-05-12T07:31:35.178881Z","shell.execute_reply":"2023-05-12T07:31:35.572695Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:35:40.878390Z","iopub.execute_input":"2023-05-12T07:35:40.879061Z","iopub.status.idle":"2023-05-12T07:35:40.907722Z","shell.execute_reply.started":"2023-05-12T07:35:40.879024Z","shell.execute_reply":"2023-05-12T07:35:40.906415Z"},"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-05-12T07:35:44.095689Z","iopub.execute_input":"2023-05-12T07:35:44.096041Z","iopub.status.idle":"2023-05-12T07:35:44.108677Z","shell.execute_reply.started":"2023-05-12T07:35:44.096013Z","shell.execute_reply":"2023-05-12T07:35:44.105764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = InkDetector().to(DEVICE)\nmodel_weights = torch.load(\"/kaggle/input/pre-trained-ink-detector/model2.pt\")\nmodel.load_state_dict(model_weights)","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:35:47.114684Z","iopub.execute_input":"2023-05-12T07:35:47.115400Z","iopub.status.idle":"2023-05-12T07:35:50.113728Z","shell.execute_reply.started":"2023-05-12T07:35:47.115356Z","shell.execute_reply":"2023-05-12T07:35:50.112731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_fragments = [test_path / fragment_name for fragment_name in test_path.iterdir()] ","metadata":{"execution":{"iopub.status.busy":"2023-05-12T07:35:50.115714Z","iopub.execute_input":"2023-05-12T07:35:50.116254Z","iopub.status.idle":"2023-05-12T07:35:50.124598Z","shell.execute_reply.started":"2023-05-12T07:35:50.116193Z","shell.execute_reply":"2023-05-12T07:35:50.123476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\n#Detect whether we are in the \"real\" evaluation by comparing the pixel values\nfrom PIL import Image\n\ndef get_num_pixels(filepath):\n    width, height = Image.open(filepath).size\n    return width*height\n\nPIXELS_B = 34523820\nPIXELS_A = 17261910\n\nCOMPETITION_RUN = False\n\nfor fragment in test_fragments:\n    \n    pixels = get_num_pixels(fragment / \"mask.png\")\n    if not pixels in [PIXELS_B, PIXELS_A]:\n        COMPETITION_RUN = True\n\nprint(\"Competition run? \", COMPETITION_RUN)\n#COMPETITION_RUN = True","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:38:37.706926Z","iopub.execute_input":"2023-05-11T21:38:37.707286Z","iopub.status.idle":"2023-05-11T21:38:37.717832Z","shell.execute_reply.started":"2023-05-11T21:38:37.707256Z","shell.execute_reply":"2023-05-11T21:38:37.716909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_images = []\nmodel.eval()\nfor test_fragment in test_fragments:\n    \n    #Remove later\n    if not COMPETITION_RUN:\n        print(\"Not in a competition run so let's not waste our time\")\n        continue\n    \n    print(\"Loading test fragment: \",test_fragment)\n    \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\n                                 , 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-05-11T21:38:38.009726Z","iopub.execute_input":"2023-05-11T21:38:38.010071Z","iopub.status.idle":"2023-05-11T21:38:38.021753Z","shell.execute_reply.started":"2023-05-11T21:38:38.010041Z","shell.execute_reply":"2023-05-11T21:38:38.020752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not COMPETITION_RUN:\n\n    pred_images = []\n    for test_fragment in test_path.iterdir():\n\n        mask_file = test_fragment / \"mask.png\"\n        mask = np.array(Image.open(mask_file).convert(\"1\"))\n\n        #Simply use mask as test-prediction\n        #But randomly select values within mask\n        random_mask = np.random.rand(*np.shape(mask))\n        pred = mask * random_mask\n\n        pred_images.append(pred)    \n\n    len(pred_images)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:38:40.111970Z","iopub.execute_input":"2023-05-11T21:38:40.112354Z","iopub.status.idle":"2023-05-11T21:38:41.054788Z","shell.execute_reply.started":"2023-05-11T21:38:40.112318Z","shell.execute_reply":"2023-05-11T21:38:41.053783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle (img, threshold):\n    flat_img = img.flatten()\n    flat_img = np.where(flat_img > threshold, 1, 0).astype(np.uint8)\n\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\n    return \" \".join(str(start) +\" \"+str(length) for start, length in zip(starts_ix, lengths))\n\n    #return \" \".join(map(str, sum(zip(starts_ix, lengths), ())))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:37:08.270206Z","iopub.execute_input":"2023-05-11T21:37:08.271182Z","iopub.status.idle":"2023-05-11T21:37:08.279225Z","shell.execute_reply.started":"2023-05-11T21:37:08.271135Z","shell.execute_reply":"2023-05-11T21:37:08.278219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"threshold = 0.4\n\nsubmission = defaultdict(list)\nfor fragment_id, fragment_name in enumerate(test_fragments):\n        print(fragment_id, fragment_name.name)\n        submission[\"Id\"].append(fragment_name.name)\n        submission[\"Predicted\"].append(rle(pred_images[fragment_id], threshold))\n\npd.DataFrame.from_dict(submission).to_csv(\"/kaggle/working/submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:37:09.613764Z","iopub.execute_input":"2023-05-11T21:37:09.614118Z","iopub.status.idle":"2023-05-11T21:37:29.090371Z","shell.execute_reply.started":"2023-05-11T21:37:09.614088Z","shell.execute_reply":"2023-05-11T21:37:29.089364Z"},"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Es kommt Vendetta\")","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:37:29.092050Z","iopub.execute_input":"2023-05-11T21:37:29.092409Z","iopub.status.idle":"2023-05-11T21:37:29.098879Z","shell.execute_reply.started":"2023-05-11T21:37:29.092377Z","shell.execute_reply":"2023-05-11T21:37:29.097967Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}