{"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":"# **Ink Detector**\nThis is a notebook explaining the Ink Detection progress [prize on Kaggle](https://www.kaggle.com/competitions/vesuvius-challenge), which is part of the larger [Vesuvius Challenge](https://scrollprize.org/).\n\nFor more background on the process of ink detection, be sure to check out [Tutorial 4: Ink Detection](https://scrollprize.org/tutorial4) on the Vesuvius Challenge website.\n\nIn this notebook we'll see how to train a simple ML model to detect ink in a papyrus fragment from a 3d x-ray scan of the fragment.\n\n\n\nFirst, initialize some variables, and let's look at a photo of the fragment. We won't use this for training, but it's useful to see.\n\nIt's an infrared photo, since the ink is better visible in infrared light.\n\n<div style=\"text-align: center;\">\n  <img src=\"https://user-images.githubusercontent.com/177461/224853397-3cf86dc2-45b4-4e7c-9ec2-28a733791a75.jpg\" alt=\"_\" style=\"height: 300px; width: 225px;\">\n</div>\n","metadata":{}},{"cell_type":"markdown","source":"## First, we've started with the Tutorial ink, to later continue with the optimization of code: \n\n- Here plotting the image of dir **\"train\"**","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport glob\nimport PIL.Image as Image\nimport torch.utils.data as data\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom tqdm import tqdm\nfrom ipywidgets import interact, fixed\nfrom segment_anything import SamPredictor, sam_model_registry\n\nPREFIX = '/kaggle/input/vesuvius-challenge-ink-detection/train/1/'\nBUFFER = 30  # Buffer size in x and y direction\nZ_START = 27  # First slice in the z direction to use\nZ_DIM = 10   # Number of slices in the z direction\nTRAINING_STEPS = 20000\nLEARNING_RATE = 0.05\nBATCH_SIZE = 24\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Paso 1: Obtener la imagen de entrada\nimage_path = PREFIX + \"ir.png\"\ninput_image = Image.open(image_path)\nplt.imshow(Image.open(PREFIX+\"ir.png\"), cmap=\"gray\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T04:39:55.279794Z","iopub.execute_input":"2023-06-14T04:39:55.28022Z","iopub.status.idle":"2023-06-14T04:40:02.454384Z","shell.execute_reply.started":"2023-06-14T04:39:55.280189Z","shell.execute_reply":"2023-06-14T04:40:02.452527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Let's load these binary images**:\n\n- **mask.png**: a mask of which pixels contain data, and which pixels we should ignore.\n- **inklabels.png**: our label data: whether a pixel contains ink or no ink (which has been hand-labeled based on the infrared photo).","metadata":{}},{"cell_type":"code","source":"# Load the 'mask.png' image and convert it to binary format ('1')\nmask = np.array(Image.open(PREFIX + \"mask.png\").convert('1'))\n\n# Load the 'inklabels.png' image and convert it to a NumPy array\n# Create a binary mask by checking if each pixel value is greater than 0\n# Convert the binary mask to a float tensor and move it to the specified device (presumably a GPU)\nlabel = torch.from_numpy(np.array(Image.open(PREFIX + \"inklabels.png\"))).gt(0).float().to(DEVICE)\n\n# Create a matplotlib figure with two subplots arranged horizontally\nfig, (ax1, ax2) = plt.subplots(1, 2)\n\n# Set the title of the first subplot to 'mask.png'\nax1.set_title(\"mask.png\")\n\n# Display the 'mask' binary image in grayscale colormap\nax1.imshow(mask, cmap='gray')\n\n# Set the title of the second subplot to 'inklabels.png'\nax2.set_title(\"inklabels.png\")\n\n# Display the 'label' binary mask (after moving it to CPU) in grayscale colormap\nax2.imshow(label.cpu(), cmap='gray')\n\n# Display the figure with the two subplots\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T04:40:30.443073Z","iopub.execute_input":"2023-06-14T04:40:30.443676Z","iopub.status.idle":"2023-06-14T04:40:37.507724Z","shell.execute_reply.started":"2023-06-14T04:40:30.443643Z","shell.execute_reply":"2023-06-14T04:40:37.506846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Next, we'll load the 3d x-ray of the fragment**. \n\nThis is represented as a .tif image stack. The image stack is an array of 16-bit grayscale images. Each image represents a \"slice\" in the z-direction, going from below the papyrus, to above the papyrus. We'll convert it to a 4D tensor of 32-bit floats. We'll also convert the pixel values to the range [0, 1].\n\nTo save memory, we'll only load the innermost slices (Z_DIM of them). Let's look at them when we're done.","metadata":{}},{"cell_type":"code","source":"# Load the 3D X-ray scan, one slice at a time\nimages = [np.array(Image.open(filename), dtype=np.float32)/65535.0 for filename in tqdm(sorted(glob.glob(PREFIX+\"surface_volume/*.tif\"))[Z_START:Z_START+Z_DIM])]\n\n# Convert the list of images into a stack tensor along the specified dimension and move it to the specified device\nimage_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n\n# Create a matplotlib figure with subplots based on the number of images\nfig, axes = plt.subplots(1, len(images), figsize=(15, 3))\n\n# Iterate over each image and its corresponding axis\nfor image, ax in zip(images, axes):\n    # Resize the image using bilinear interpolation and convert it to a float32 array\n    resized_image = np.array(Image.fromarray(image).resize((image.shape[1]//20, image.shape[0]//20)), dtype=np.float32)\n    \n    # Display the resized image in grayscale colormap\n    ax.imshow(resized_image, cmap='gray')\n    \n    # Remove the ticks on the x and y axes\n    ax.set_xticks([])\n    ax.set_yticks([])\n\n# Adjust the spacing between subplots\nfig.tight_layout()\n\n# Display the figure with the subplots\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:03:45.613576Z","iopub.execute_input":"2023-06-14T05:03:45.613963Z","iopub.status.idle":"2023-06-14T05:04:10.809053Z","shell.execute_reply.started":"2023-06-14T05:03:45.613933Z","shell.execute_reply":"2023-06-14T05:04:10.80822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Can you see the ink in these slices of the 3d x-ray scan..? **Neither can we**.\n\n## **Now we'll create a dataset of subvolumes**. \nWe use a small rectangle around the letter \"P\" for our evaluation, and we'll exclude those pixels from the training set. \n\n(**It's actually a Greek letter \"rho\", which looks similar to our \"P\".**)","metadata":{}},{"cell_type":"code","source":"# Define the coordinates of the rectangle to create a subvolume dataset\nrect = (1100, 3500, 700, 950)\n\n# Create a matplotlib figure and axis\nfig, ax = plt.subplots()\n\n# Display the binary mask stored in 'label' (after moving it to CPU) in grayscale colormap\nax.imshow(label.cpu())\n\n# Create a patch (rectangle) with specified dimensions and appearance\npatch = patches.Rectangle((rect[0], rect[1]), rect[2], rect[3], linewidth=2, edgecolor='r', facecolor='none')\n\n# Add the patch to the axis\nax.add_patch(patch)\n\n# Display the figure with the binary mask and the highlighted rectangle\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:08:26.489137Z","iopub.execute_input":"2023-06-14T05:08:26.489508Z","iopub.status.idle":"2023-06-14T05:08:28.323581Z","shell.execute_reply.started":"2023-06-14T05:08:26.489479Z","shell.execute_reply":"2023-06-14T05:08:28.322701Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Now we'll define a PyTorch dataset and (super simple) model.**","metadata":{}},{"cell_type":"code","source":"# Define a PyTorch dataset for the subvolumes\nclass SubvolumeDataset(data.Dataset):\n    def __init__(self, image_stack, label, pixels):\n        self.image_stack = image_stack\n        self.label = label\n        self.pixels = pixels\n    \n    def __len__(self):\n        return len(self.pixels)\n    \n    def __getitem__(self, index):\n        # Get the coordinates of the pixel at the specified index\n        y, x = self.pixels[index]\n        \n        # Extract the subvolume around the pixel from the image stack\n        subvolume = self.image_stack[:, y-BUFFER:y+BUFFER+1, x-BUFFER:x+BUFFER+1].view(1, Z_DIM, BUFFER*2+1, BUFFER*2+1)\n        \n        # Get the ink label for the pixel\n        inklabel = self.label[y, x].view(1)\n        \n        return subvolume, inklabel\n\n# Define the model architecture using PyTorch's Sequential API\nmodel = nn.Sequential(\n    nn.Conv3d(1, 16, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(16, 32, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(32, 64, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Flatten(start_dim=1),\n    nn.LazyLinear(128), nn.ReLU(),\n    nn.LazyLinear(1), nn.Sigmoid()\n).to(DEVICE)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:10:54.487605Z","iopub.execute_input":"2023-06-14T05:10:54.488005Z","iopub.status.idle":"2023-06-14T05:10:54.516807Z","shell.execute_reply.started":"2023-06-14T05:10:54.487975Z","shell.execute_reply":"2023-06-14T05:10:54.515925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Now we'll train the model. Conceptually it looks like this:**\n<div style=\"text-align: center;\">\n  <img src=\"https://user-images.githubusercontent.com/22727759/224853385-ed190d89-f466-469c-82a9-499881759d57.gif\" alt=\"_\" style=\"height: 275px; width: 585px;\">\n</div>","metadata":{}},{"cell_type":"markdown","source":"## **This typically takes about 10 minutes.**","metadata":{}},{"cell_type":"code","source":"# Generate pixel lists for training and validation sets\nprint(\"Generating pixel lists...\")\n\n# Create a Boolean array 'not_border' of the same shape as the bitmask, initially all True\nnot_border = np.zeros(mask.shape, dtype=bool)\nnot_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n\n# Apply the bitmask to the 'mask' array, keeping only the pixels inside the mask but not on the border\narr_mask = np.array(mask) * not_border\n\n# Create a Boolean array 'inside_rect' and set it to True for pixels within the specified rectangle\ninside_rect = np.zeros(mask.shape, dtype=bool) * arr_mask\ninside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = True\n\n# Create a Boolean array 'outside_rect' and set it to True for pixels outside the specified rectangle\noutside_rect = np.ones(mask.shape, dtype=bool) * arr_mask\noutside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = False\n\n# Get the pixel coordinates inside and outside the rectangle\npixels_inside_rect = np.argwhere(inside_rect)\npixels_outside_rect = np.argwhere(outside_rect)\n\nprint(\"Training...\")\n\n# Create the training dataset using the SubvolumeDataset class with pixels outside the rectangle\ntrain_dataset = SubvolumeDataset(image_stack, label, pixels_outside_rect)\n\n# Create a data loader for the training dataset\ntrain_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\n# Define the loss function and optimizer\ncriterion = nn.BCELoss()\noptimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE)\n\n# Create a learning rate scheduler\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\n\n# Set the model to training mode\nmodel.train()\n\n# Iterate over the training loader for the specified number of training steps\nfor i, (subvolumes, inklabels) in tqdm(enumerate(train_loader), total=TRAINING_STEPS):\n    if i >= TRAINING_STEPS:\n        break\n    \n    # Zero the gradients\n    optimizer.zero_grad()\n    \n    # Forward pass through the model\n    outputs = model(subvolumes.to(DEVICE))\n    \n    # Compute the loss\n    loss = criterion(outputs, inklabels.to(DEVICE))\n    \n    # Backward pass and optimization\n    loss.backward()\n    optimizer.step()\n    \n    # Adjust the learning rate\n    scheduler.step()\n    \n# Calculate the running loss by accumulating the loss value\n# running_loss += loss.item()\n\n# Print the average loss every 3000 iterations\n# if i % 3000 == 3000-1:\n#     print(\"Loss:\", running_loss / 3000)\n#     running_loss = 0.0","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:18:22.691595Z","iopub.execute_input":"2023-06-14T05:18:22.691989Z","iopub.status.idle":"2023-06-14T05:22:58.544829Z","shell.execute_reply.started":"2023-06-14T05:18:22.691959Z","shell.execute_reply":"2023-06-14T05:22:58.543739Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Finally, we'll generate a prediction image**. \n\nWe'll use the model to predict the presence of ink for each pixel in our rectangle **(the val set)**. Conceptually it looks like this:\n\n<div style=\"text-align: center;\">\n  <img src=\"https://user-images.githubusercontent.com/22727759/224853379-09ae991e-02be-4ecc-a652-313165b3005c.gif\" alt=\"_\" style=\"height: 275px; width: 585px;\">\n</div>\n\n---\n\nThis should take about a minutes.\n\n**Remember that the model has never seen the label data within the rectangle before!**\n\nWe'll plot it side-by-side with the label image. Are you able to recognize the letter **\"P\"** in it?","metadata":{}},{"cell_type":"code","source":"# Create an evaluation dataset using the SubvolumeDataset class\neval_dataset = SubvolumeDataset(image_stack, label, pixels_inside_rect)\n\n# Create a data loader for the evaluation dataset\neval_loader = data.DataLoader(eval_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# Create a tensor 'output' with the same size as the label and initialize it with zeros\noutput = torch.zeros_like(label).float()\n\n# Set the model to evaluation mode\nmodel.eval()\n\n# Disable gradient calculations since we're in evaluation mode\nwith torch.no_grad():\n    # Iterate over the evaluation loader\n    for i, (subvolumes, _) in enumerate(tqdm(eval_loader)):\n        # Pass the subvolumes through the model and get the predicted values\n        predictions = model(subvolumes.to(DEVICE))\n        \n        # Iterate over the predicted values and their corresponding indices\n        for j, value in enumerate(predictions):\n            # Assign the predicted value to the corresponding pixel in the 'output' tensor\n            output[tuple(pixels_inside_rect[i*BATCH_SIZE+j])] = value\n\n# Create a matplotlib figure with two subplots\nfig, (ax1, ax2) = plt.subplots(1, 2)\n\n# Display the 'output' tensor (after moving it to CPU) in grayscale colormap\nax1.imshow(output.cpu(), cmap='gray')\n\n# Display the original label (after moving it to CPU) in grayscale colormap\nax2.imshow(label.cpu(), cmap='gray')\n\n# Show the figure with the two subplots\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:26:34.925963Z","iopub.execute_input":"2023-06-14T05:26:34.926395Z","iopub.status.idle":"2023-06-14T05:28:26.897715Z","shell.execute_reply.started":"2023-06-14T05:26:34.926363Z","shell.execute_reply":"2023-06-14T05:28:26.896654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Since our output has to be binary, we have to choose a threshold, say 40% confidence.**","metadata":{}},{"cell_type":"code","source":"# Define the threshold for the binary output\nTHRESHOLD = 0.4\n\n# Create a matplotlib figure with two subplots\nfig, (ax1, ax2) = plt.subplots(1, 2)\n\n# Display the binary output by applying the threshold to the 'output' tensor (after moving it to CPU) in grayscale colormap\nax1.imshow(output.gt(THRESHOLD).cpu(), cmap='gray')\n\n# Display the original label (after moving it to CPU) in grayscale colormap\nax2.imshow(label.cpu(), cmap='gray')\n\n# Show the figure with the two subplots\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:38:01.524605Z","iopub.execute_input":"2023-06-14T05:38:01.525005Z","iopub.status.idle":"2023-06-14T05:38:04.555222Z","shell.execute_reply.started":"2023-06-14T05:38:01.524973Z","shell.execute_reply":"2023-06-14T05:38:04.554158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Finally, Kaggle expects a run-length encoded submit.csv file, so let's generate that.**\n\nLet's skip this step to give way to explain the suggestions and different types of models to make the comparisons.","metadata":{}},{"cell_type":"code","source":"# Adapted from https://www.kaggle.com/code/stainsby/fast-tested-rle/notebook\n# and https://www.kaggle.com/code/kotaiizuka/faster-rle/notebook\n# and https://www.kaggle.com/code/jpposma/vesuvius-challenge-ink-detection-tutorial\n\n# Define a function to convert the binary output to the required run-length encoding (RLE) format\ndef rle(output):\n    # Convert the binary output to an array of 1s and 0s using the threshold\n    pixels = np.where(output.flatten().cpu() > THRESHOLD, 1, 0).astype(np.uint8)\n    \n    # Set the first and last pixels to 0 (assumed background)\n    pixels[0] = 0\n    pixels[-1] = 0\n    \n    # Find the runs (consecutive occurrences) of 1s and their lengths\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] = runs[1::2] - runs[:-1:2]  # Compute the lengths of the runs\n    \n    # Convert the runs to a string in the required format\n    return ' '.join(str(x) for x in runs)\n\n# Apply the rle function to the output tensor to get the RLE-encoded output\nrle_output = rle(output)\n\n# Output the RLE-encoded output in the required format to a file named 'submission_tutorial.csv'\nprint(\"Id,Predicted\\na,\" + rle_output + \"\\nb,\" + rle_output, file=open('submission_tutorial.csv', 'w'))\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T05:48:16.198412Z","iopub.execute_input":"2023-06-14T05:48:16.198775Z","iopub.status.idle":"2023-06-14T05:48:16.936586Z","shell.execute_reply.started":"2023-06-14T05:48:16.198746Z","shell.execute_reply":"2023-06-14T05:48:16.935639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = pd.read_csv('/kaggle/working/submission_tutorial.csv')\ndata","metadata":{"execution":{"iopub.status.busy":"2023-06-14T06:04:01.771724Z","iopub.execute_input":"2023-06-14T06:04:01.772725Z","iopub.status.idle":"2023-06-14T06:04:01.824557Z","shell.execute_reply.started":"2023-06-14T06:04:01.772691Z","shell.execute_reply":"2023-06-14T06:04:01.823493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Example submission**","metadata":{}},{"cell_type":"markdown","source":"## **Import libraries**:","metadata":{}},{"cell_type":"code","source":"import os  # Provides functions for interacting with the operating system\nimport gc  # Performs garbage collection to free up memory\nimport glob  # Helps in finding files using pattern matching\nimport json  # Handles JSON data\nfrom collections import defaultdict  # Implements a dictionary with default values\nimport multiprocessing as mp  # Supports multiprocessing\nfrom pathlib import Path  # Represents file system paths\nfrom types import SimpleNamespace  # Creates simple objects with named attributes\nfrom typing import Dict, List, Optional, Tuple  # Provides type hints for variables and functions\nimport warnings  # Controls warning messages\n\nimport matplotlib.pyplot as plt  # Produces plots and visualizations\nimport matplotlib.patches as patches  # Represents geometric shapes in plots\nimport numpy as np  # Handles numerical operations and array manipulation\nimport pandas as pd  # Manipulates and analyzes data in tabular form\nimport PIL.Image as Image  # Deals with image-related operations\nfrom sklearn.metrics import fbeta_score  # Computes the F-beta score for binary classification\nfrom sklearn.exceptions import UndefinedMetricWarning  # Handles warnings for undefined metrics\nimport torch  # Provides support for deep learning with tensors\nimport torch.nn as nn  # Implements neural network modules and functions\nimport torch.optim as optim  # Optimizes parameters for training neural networks\nimport torch.utils.data as thd  # Provides tools for working with datasets and data loaders\nfrom tqdm import tqdm  # Displays progress bars during iterations","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:29:09.43446Z","iopub.execute_input":"2023-06-15T03:29:09.434929Z","iopub.status.idle":"2023-06-15T03:29:12.919568Z","shell.execute_reply.started":"2023-06-15T03:29:09.434894Z","shell.execute_reply":"2023-06-15T03:29:12.918639Z"},"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            # Load image stack\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            # Collect valid pixels\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\n            # Encode fragment ID in pixels\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            # Load inklabels if enabled\n            if load_inklabels:\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 the subvolume exceeds the image dimensions\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            # Create a subvolume array and fill it with image data\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            # Extract the subvolume from the image stack\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\n        # Get the corresponding inklabel for the subvolume\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        # Plot the label image and highlight the subvolume region\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-06-15T03:29:17.9777Z","iopub.execute_input":"2023-06-15T03:29:17.978267Z","iopub.status.idle":"2023-06-15T03:29:18.006498Z","shell.execute_reply.started":"2023-06-15T03:29:17.978237Z","shell.execute_reply":"2023-06-15T03:29:18.00553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_path = Path(\"/kaggle/input/vesuvius-challenge-ink-detection/\")\ntrain_path = base_path / \"train\"\n\n# Get the names of all fragments in the train directory\nall_fragments = sorted([f.name for f in train_path.iterdir()])\nprint(\"All fragments:\", all_fragments)\n\n# Select a subset of fragments to use for training\n# In this example, we choose the fragment with name \"1\"\ntrain_fragments = [train_path / fragment_name for fragment_name in [\"1\"]]\n\n# The variable \"train_fragments\" now contains the paths to the selected fragments","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:29:19.339695Z","iopub.execute_input":"2023-06-15T03:29:19.340313Z","iopub.status.idle":"2023-06-15T03:29:19.351106Z","shell.execute_reply.started":"2023-06-15T03:29:19.340283Z","shell.execute_reply":"2023-06-15T03:29:19.350083Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Measure the execution time\ntrain_dset = SubvolumeDataset(fragments=train_fragments, voxel_shape=(48, 64, 64), filter_edge_pixels=True)\n# Create an instance of the SubvolumeDataset class with the specified arguments\n# - fragments: train_fragments, which contains the selected fragment paths\n# - voxel_shape: (48, 64, 64), specifying the shape of the subvolumes\n# - filter_edge_pixels: True, indicating that edge pixels should be filtered\n\nprint(\"Num items (pixels)\", len(train_dset))\n# Print the number of items (pixels) in the train dataset using the len() function","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:29:28.891049Z","iopub.execute_input":"2023-06-15T03:29:28.891405Z","iopub.status.idle":"2023-06-15T03:31:00.920395Z","shell.execute_reply.started":"2023-06-15T03:29:28.891377Z","shell.execute_reply":"2023-06-15T03:31:00.919402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Sanity Check**","metadata":{}},{"cell_type":"code","source":"# Sanity check\nindex = 6136130\n# Set the index for the sanity check\n\ntrain_dset.plot_label(index, figsize=(16, 10))\n# Call the plot_label() method of the train_dset object to visualize the label at the specified index\n# Pass the figsize argument to control the size of the figure","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:31:00.922217Z","iopub.execute_input":"2023-06-15T03:31:00.923146Z","iopub.status.idle":"2023-06-15T03:31:03.75619Z","shell.execute_reply.started":"2023-06-15T03:31:00.923111Z","shell.execute_reply":"2023-06-15T03:31:03.755276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 32\n# Set the batch size for the data loader\n\ntrain_loader = thd.DataLoader(train_dset, batch_size=BATCH_SIZE, shuffle=True)\n# Create a data loader for the training dataset\n# Pass the train_dset object as the dataset\n# Set the batch_size to BATCH_SIZE\n# Set shuffle=True to shuffle the data during training\n\nprint(\"Num batches:\", len(train_loader))\n# Print the number of batches in the train_loader\n# This represents the number of iterations required to cover the entire training dataset","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:31:03.757186Z","iopub.execute_input":"2023-06-15T03:31:03.757516Z","iopub.status.idle":"2023-06-15T03:31:03.765105Z","shell.execute_reply.started":"2023-06-15T03:31:03.757486Z","shell.execute_reply":"2023-06-15T03:31:03.764052Z"},"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\")\n# Set the device to CUDA if available, otherwise use CPU","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:31:03.767712Z","iopub.execute_input":"2023-06-15T03:31:03.768437Z","iopub.status.idle":"2023-06-15T03:31:03.834869Z","shell.execute_reply.started":"2023-06-15T03:31:03.768348Z","shell.execute_reply":"2023-06-15T03:31:03.833673Z"},"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]  # Number of filters in each convolutional layer\n        paddings = [1, 1, 1]  # Padding size in each convolutional layer\n        kernel_sizes = [3, 3, 3]  # Kernel size in each convolutional layer\n        strides = [2, 2, 2]  # Stride size in each convolutional layer\n\n        layers = []  # List to store the layers of the model\n        in_channels = 1  # Number of input channels\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  # Update the number of input channels for the next layer\n        layers.append(nn.AdaptiveAvgPool3d(1))  # Adaptive average pooling layer\n        layers.append(nn.Flatten())  # Flatten the features\n\n        self.encoder = nn.Sequential(*layers)  # Encoder part of the model\n        self.decoder = nn.Sequential(\n            nn.Linear(in_channels, 128),  # Fully connected layer\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 128),  # Fully connected layer\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 1)  # Output layer\n        )\n\n    def forward(self, x):\n        features = self.encoder(x)  # Pass the input through the encoder\n        return self.decoder(features)  # Pass the features through the decoder\n","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:33:11.714799Z","iopub.execute_input":"2023-06-15T03:33:11.715721Z","iopub.status.idle":"2023-06-15T03:33:11.726862Z","shell.execute_reply.started":"2023-06-15T03:33:11.715678Z","shell.execute_reply":"2023-06-15T03:33:11.725851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of the InkDetector model and move it to the specified device\nmodel = InkDetector().to(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:33:16.561Z","iopub.execute_input":"2023-06-15T03:33:16.561362Z","iopub.status.idle":"2023-06-15T03:33:19.514726Z","shell.execute_reply.started":"2023-06-15T03:33:16.561333Z","shell.execute_reply":"2023-06-15T03:33:19.513786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Train**:","metadata":{}},{"cell_type":"code","source":"# Define the number of training steps and learning rate\nTRAINING_STEPS = 30000\nLEARNING_RATE = 1e-3\n\n# Set TRAIN_RUN to True to enable training\nTRAIN_RUN = True  # To avoid re-running when saving the notebook","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:33:19.690036Z","iopub.execute_input":"2023-06-15T03:33:19.6907Z","iopub.status.idle":"2023-06-15T03:33:19.697057Z","shell.execute_reply.started":"2023-06-15T03:33:19.690667Z","shell.execute_reply":"2023-06-15T03:33:19.694004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Ignore UndefinedMetricWarning\nwarnings.simplefilter('ignore', UndefinedMetricWarning)","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:33:25.417424Z","iopub.execute_input":"2023-06-15T03:33:25.418086Z","iopub.status.idle":"2023-06-15T03:33:25.422409Z","shell.execute_reply.started":"2023-06-15T03:33:25.418053Z","shell.execute_reply":"2023-06-15T03:33:25.42142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if TRAIN_RUN:\n    # Training mode\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\n    # Initialize running metrics\n    running_loss = 0.0\n    running_accuracy = 0.0\n    running_fbeta = 0.0\n    denom = 0\n\n    # Create a progress bar using tqdm\n    pbar = tqdm(enumerate(train_loader), total=TRAINING_STEPS)\n\n    # Training loop\n    for i, (subvolumes, inklabels) in pbar:\n        if i >= TRAINING_STEPS:\n            break\n\n        # Zero the gradients\n        optimizer.zero_grad()\n\n        # Forward pass\n        outputs = model(subvolumes.to(DEVICE))\n\n        # Calculate the loss\n        loss = criterion(outputs, inklabels.to(DEVICE))\n\n        # Backward pass and optimization\n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n\n        # Compute predictions and metrics\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\n        # Update the progress bar with the current metrics\n        pbar.set_postfix({\"Loss\": running_loss / denom, \"Accuracy\": running_accuracy / denom, \"Fbeta@0.5\": running_fbeta / denom})\n\n        # Reset the running metrics every 500 steps\n        if (i + 1) % 500 == 0:\n            running_loss = 0.\n            running_accuracy = 0.\n            running_fbeta = 0.\n            denom = 0\n\n    # Save the trained model\n    torch.save(model.state_dict(), \"/kaggle/working/model.pt\")\n\nelse:\n    # Load pre-trained model weights\n    model_weights = torch.load(\"/kaggle/working/model.pt\")\n    model.load_state_dict(model_weights)","metadata":{"execution":{"iopub.status.busy":"2023-06-15T03:36:30.763834Z","iopub.execute_input":"2023-06-15T03:36:30.764275Z","iopub.status.idle":"2023-06-15T04:00:30.016877Z","shell.execute_reply.started":"2023-06-15T03:36:30.764239Z","shell.execute_reply":"2023-06-15T04:00:30.015796Z"},"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-06-15T04:00:30.018976Z","iopub.execute_input":"2023-06-15T04:00:30.019344Z","iopub.status.idle":"2023-06-15T04:00:30.571231Z","shell.execute_reply.started":"2023-06-15T04:00:30.01931Z","shell.execute_reply":"2023-06-15T04:00:30.570297Z"},"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-06-15T04:00:30.572525Z","iopub.execute_input":"2023-06-15T04:00:30.573628Z","iopub.status.idle":"2023-06-15T04:00:30.595525Z","shell.execute_reply.started":"2023-06-15T04:00:30.573568Z","shell.execute_reply":"2023-06-15T04:00:30.594137Z"},"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-06-15T04:22:47.325268Z","iopub.execute_input":"2023-06-15T04:22:47.325823Z"},"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)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.DataFrame.from_dict(submission)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip install git+https://github.com/facebookresearch/segment-anything.git","metadata":{"execution":{"iopub.status.busy":"2023-06-15T02:05:47.311341Z","iopub.status.idle":"2023-06-15T02:05:47.311838Z","shell.execute_reply.started":"2023-06-15T02:05:47.311555Z","shell.execute_reply":"2023-06-15T02:05:47.311576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}