{"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":"## Purpose Statement","metadata":{}},{"cell_type":"markdown","source":"The purpose of this notebook is to replicate the original \"Vesuvius Challenge: Ink Detection tutorial\" notebook (https://www.kaggle.com/code/jpposma/vesuvius-challenge-ink-detection-tutorial) in tensorflow. The hope is that you can read the two notebocks block by block and see a parallel implementation. Thanks to Viktor Fairuschin for his implementation of the tutorial (https://www.kaggle.com/code/viktorfairuschin/ink-detektion-tutorial-with-tensorflow), which helped me figure out why my training loop was taking so long.","metadata":{}},{"cell_type":"markdown","source":"# Tensorflow Implementation","metadata":{}},{"cell_type":"markdown","source":"This 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<img src=\"https://user-images.githubusercontent.com/177461/224853397-3cf86dc2-45b4-4e7c-9ec2-28a733791a75.jpg\" width=\"200\"/>\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.","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow import keras\nimport math\nimport tensorflow.keras.layers as tfl\nimport numpy as np\nimport glob\nimport PIL.Image as Image\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom tqdm import tqdm\nfrom ipywidgets import interact, fixed\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 = 30000\nLEARNING_RATE = 0.03\nBATCH_SIZE = 32\nDEVICE = tf.distribute.MirroredStrategy()\n\nplt.imshow(Image.open(PREFIX+\"ir.png\"), cmap=\"gray\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-19T15:20:58.065766Z","iopub.execute_input":"2023-04-19T15:20:58.066391Z","iopub.status.idle":"2023-04-19T15:21:12.503533Z","shell.execute_reply.started":"2023-04-19T15:20:58.066334Z","shell.execute_reply":"2023-04-19T15:21:12.500591Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's load these binary images:\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":"mask = np.array(Image.open(PREFIX+\"mask.png\").convert('1'))\nlabel = tf.convert_to_tensor(np.array(Image.open(PREFIX+\"inklabels.png\")), dtype = tf.dtypes.bool)\nfig, (ax1, ax2) = plt.subplots(1, 2)\nax1.set_title(\"mask.png\")\nax1.imshow(mask, cmap='gray')\nax2.set_title(\"inklabels.png\")\nax2.imshow(label.cpu(), cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:21:12.505544Z","iopub.execute_input":"2023-04-19T15:21:12.506169Z","iopub.status.idle":"2023-04-19T15:21:15.251969Z","shell.execute_reply.started":"2023-04-19T15:21:12.506129Z","shell.execute_reply":"2023-04-19T15:21:15.250925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Next, we'll load the 3d x-ray of the fragment. This 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])]\nimage_stack  = tf.stack(images, axis=0)\nimage_stack = tf.expand_dims(image_stack, axis=-1)  # Adding channels axis now as opposed to later\n\nfig, axes = plt.subplots(1, len(images), figsize=(15, 3))\nfor image, ax in zip(images, axes):\n  ax.imshow(np.array(Image.fromarray(image).resize((image.shape[1]//20, image.shape[0]//20)), dtype=np.float32), cmap='gray')\n  ax.set_xticks([]); ax.set_yticks([])\nfig.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:21:15.253689Z","iopub.execute_input":"2023-04-19T15:21:15.254339Z","iopub.status.idle":"2023-04-19T15:22:15.270887Z","shell.execute_reply.started":"2023-04-19T15:21:15.254296Z","shell.execute_reply":"2023-04-19T15:22:15.269870Z"},"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\nNow we'll create a dataset of subvolumes. We use a small rectangle around the letter \"P\" for our evaluation, and we'll exclude those pixels from the training set. (It's actually a Greek letter \"rho\", which looks similar to our \"P\".)","metadata":{}},{"cell_type":"code","source":"rect = (1100, 3500, 700, 950)\nfig, ax = plt.subplots()\nax.imshow(label.cpu())\npatch = patches.Rectangle((rect[0], rect[1]), rect[2], rect[3], linewidth=2, edgecolor='r', facecolor='none')\nax.add_patch(patch)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:22:15.273696Z","iopub.execute_input":"2023-04-19T15:22:15.274466Z","iopub.status.idle":"2023-04-19T15:22:16.703793Z","shell.execute_reply.started":"2023-04-19T15:22:15.274425Z","shell.execute_reply":"2023-04-19T15:22:16.702690Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now define a tensorflow dataset and simple model","metadata":{}},{"cell_type":"code","source":"class SubvolumeDataset(tf.keras.utils.Sequence):\n    '''\n    This subclass carries out both the SubvolumeDataset(data.Dataset) and data.DataLoader steps at the same time.\n    Therefore, the logic for both must be implemented here. Namely, this involves batching and shuffling the input data.\n    '''\n    def __init__(self, image_stack, label, pixels, batch_size, shuffle = True):\n        super(SubvolumeDataset, self).__init__()\n        self.image_stack = image_stack\n        self.label = label\n        self.pixels = pixels\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        self.on_epoch_end()\n        self.length = int(np.ceil( len(self.pixels) / self.batch_size ))\n\n        if shuffle == True:\n            tf.random.shuffle(self.pixels, seed=10)\n\n    def on_epoch_end(self):\n        if self.shuffle == True:\n            tf.random.shuffle(self.pixels, seed=10)\n\n    def __len__(self):\n        # Denotes the number of batches per epoch\n        return self.length\n\n    def __getitem__(self, index):\n        # Returns the batch identified by 'index'\n        if index < self.length:\n            # Handles everything prior to the last batch\n            pixels = self.pixels[index * self.batch_size : (index+1) * self.batch_size]\n        else:\n            # Handles the last batch\n            pixels = self.pixels[index * self.batch_size:]\n\n        subvolumes = []\n        inklabels  = []\n        for y, x in pixels:\n\n            subvolume = self.image_stack[:, y-BUFFER:y+BUFFER+1, x-BUFFER:x+BUFFER+1]\n            subvolumes.append(subvolume)\n\n            inklabels.append(self.label[y, x])\n\n        # Stack the elements of the batch\n        subvolumes = tf.stack(subvolumes, axis = 0)\n        \n        # Stack the labels for each element in the batch\n        inklabels = tf.convert_to_tensor(inklabels)\n\n        return subvolumes, inklabels\n\n\nwith DEVICE.scope(): # Similar to '.to(DEVICE)'\n\n    model = tf.keras.models.Sequential([\n        tfl.Conv3D(16, 3, strides=1, padding='same', input_shape=(Z_DIM, BUFFER*2+1, BUFFER*2+1, 1)),\n        tfl.MaxPool3D(strides=2),\n        tfl.Conv3D(32, 3, strides=1, padding='same'), \n        tfl.MaxPool3D(strides=2),\n        tfl.Conv3D(64, 3, strides=1, padding='same'), \n        tfl.MaxPool3D(strides=2),\n        tfl.Flatten(),\n        tfl.Dense(128, activation='relu'),\n        tfl.Dense(1, activation='sigmoid'),\n    ])","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:22:16.707986Z","iopub.execute_input":"2023-04-19T15:22:16.708284Z","iopub.status.idle":"2023-04-19T15:22:16.861351Z","shell.execute_reply.started":"2023-04-19T15:22:16.708254Z","shell.execute_reply":"2023-04-19T15:22:16.860424Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now we'll train the model. Conceptually it looks like this:\n\n<a href=\"https://user-images.githubusercontent.com/22727759/224853655-3fad9edb-c798-452e-94d0-f74efe71c08e.mp4\"><img src=\"https://user-images.githubusercontent.com/22727759/224853385-ed190d89-f466-469c-82a9-499881759d57.gif\"/></a>\n\nThis typically takes about 10 minutes.","metadata":{}},{"cell_type":"code","source":"# Thanks to Viktor for pointing me to avanwyk's implementation of the OneCycleScheduler!\nimport logging\n\nlogging.getLogger('tensorflow').setLevel(logging.ERROR)\n\nfrom tensorflow.keras.callbacks import Callback\n\nclass CosineAnnealer:\n    \n    def __init__(self, start, end, steps):\n        self.start = start\n        self.end = end\n        self.steps = steps\n        self.n = 0\n        \n    def step(self):\n        self.n += 1\n        cos = np.cos(np.pi * (self.n / self.steps)) + 1\n        return self.end + (self.start - self.end) / 2. * cos\n\n\nclass OneCycleScheduler(Callback):\n    \"\"\" \n    https://www.kaggle.com/code/avanwyk/tf2-super-convergence-with-the-1cycle-policy/notebook\n    \"\"\"\n\n    def __init__(self, lr_max, steps, mom_min=0.85, mom_max=0.95, phase_1_pct=0.3, div_factor=25.):\n        super(OneCycleScheduler, self).__init__()\n        lr_min = lr_max / div_factor\n        final_lr = lr_max / (div_factor * 1e4)\n        phase_1_steps = steps * phase_1_pct\n        phase_2_steps = steps - phase_1_steps\n        \n        self.phase_1_steps = phase_1_steps\n        self.phase_2_steps = phase_2_steps\n        self.phase = 0\n        self.step = 0\n        \n        self.phases = [[CosineAnnealer(lr_min, lr_max, phase_1_steps), CosineAnnealer(mom_max, mom_min, phase_1_steps)], \n                 [CosineAnnealer(lr_max, final_lr, phase_2_steps), CosineAnnealer(mom_min, mom_max, phase_2_steps)]]\n        \n        self.lrs = []\n        self.moms = []\n\n    def on_train_begin(self, logs=None):\n        self.phase = 0\n        self.step = 0\n\n        self.set_lr(self.lr_schedule().start)\n        self.set_momentum(self.mom_schedule().start)\n        \n    def on_train_batch_begin(self, batch, logs=None):\n        self.lrs.append(self.get_lr())\n        self.moms.append(self.get_momentum())\n\n    def on_train_batch_end(self, batch, logs=None):\n        self.step += 1\n        if self.step >= self.phase_1_steps:\n            self.phase = 1\n            \n        self.set_lr(self.lr_schedule().step())\n        self.set_momentum(self.mom_schedule().step())\n        \n    def get_lr(self):\n        try:\n            return tf.keras.backend.get_value(self.model.optimizer.lr)\n        except AttributeError:\n            return None\n        \n    def get_momentum(self):\n        try:\n            return tf.keras.backend.get_value(self.model.optimizer.momentum)\n        except AttributeError:\n            return None\n        \n    def set_lr(self, lr):\n        try:\n            tf.keras.backend.set_value(self.model.optimizer.lr, lr)\n        except AttributeError:\n            pass # ignore\n        \n    def set_momentum(self, mom):\n        try:\n            tf.keras.backend.set_value(self.model.optimizer.momentum, mom)\n        except AttributeError:\n            pass # ignore\n\n    def lr_schedule(self):\n        return self.phases[self.phase][0]\n    \n    def mom_schedule(self):\n        return self.phases[self.phase][1]\n    \n    def plot(self):\n        ax = plt.subplot(1, 2, 1)\n        ax.plot(self.lrs)\n        ax.set_title('Learning Rate')\n        ax = plt.subplot(1, 2, 2)\n        ax.plot(self.moms)\n        ax.set_title('Momentum')","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:22:16.862897Z","iopub.execute_input":"2023-04-19T15:22:16.863256Z","iopub.status.idle":"2023-04-19T15:22:16.883179Z","shell.execute_reply.started":"2023-04-19T15:22:16.863220Z","shell.execute_reply":"2023-04-19T15:22:16.882180Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Generating pixel lists...\")\n# Split our dataset into train and val. The pixels inside the rect are the \n# val set, and the pixels outside the rect are the train set.\n# Adapted from https://www.kaggle.com/code/jamesdavey/100x-faster-pixel-coordinate-generator-1s-runtime\n# Create a Boolean array 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\narr_mask = np.array(mask) * not_border\ninside_rect = np.zeros(mask.shape, dtype=bool) * arr_mask\n# Sets all indexes with inside_rect array to True\ninside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = True\n# Set the pixels within the inside_rect to False\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\npixels_inside_rect = np.argwhere(inside_rect)\npixels_outside_rect = np.argwhere(outside_rect)\n\n# Training is going a bit slower than I would like ~ 17.5 min\nprint('Training...')\ntrain_loader = SubvolumeDataset(image_stack, label, pixels_outside_rect, BATCH_SIZE, shuffle = True)\ncriterion = tf.keras.losses.BinaryCrossentropy(from_logits=False)\noptimizer = tf.keras.optimizers.SGD(lr=LEARNING_RATE)\nscheduler = OneCycleScheduler(lr_max=LEARNING_RATE, steps=TRAINING_STEPS)\nwith DEVICE.scope():\n    model.compile(loss=criterion, optimizer=optimizer, metrics=['accuracy'])\nhistory = model.fit(train_loader, epochs=1, steps_per_epoch=TRAINING_STEPS, callbacks=[scheduler])","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:22:16.884434Z","iopub.execute_input":"2023-04-19T15:22:16.885293Z","iopub.status.idle":"2023-04-19T15:39:55.168796Z","shell.execute_reply.started":"2023-04-19T15:22:16.885256Z","shell.execute_reply":"2023-04-19T15:39:55.167668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This loop is also going slower than I would like ~ 15 min. Also, performance is slightly worse.\neval_dataset = SubvolumeDataset(image_stack, label, pixels_inside_rect, BATCH_SIZE, shuffle = False)\noutput = np.zeros_like(label, dtype='float32')\nfor i, (subvolumes, _) in enumerate(tqdm(eval_dataset)):\n    for j, value in enumerate(model(subvolumes, training=False)):\n        #output[tuple(pixels_inside_rect[i*BATCH_SIZE+j])].assign(value)\n        #output = outputs\n        output[tuple(pixels_inside_rect[i*BATCH_SIZE+j])] = value\n        \nfig, (ax1, ax2) = plt.subplots(1, 2)\nax1.imshow(output, cmap='gray')\nax2.imshow(label, cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T15:42:11.279211Z","iopub.execute_input":"2023-04-19T15:42:11.279928Z","iopub.status.idle":"2023-04-19T15:57:17.768740Z","shell.execute_reply.started":"2023-04-19T15:42:11.279888Z","shell.execute_reply":"2023-04-19T15:57:17.767202Z"},"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.\n\n","metadata":{}},{"cell_type":"code","source":"THRESHOLD = 0.4\nfig, (ax1, ax2) = plt.subplots(1, 2)\nax1.imshow(np.where(output > THRESHOLD, 1, 0), cmap='gray')\nax2.imshow(label.cpu(), cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:19:05.173584Z","iopub.execute_input":"2023-04-19T16:19:05.174612Z","iopub.status.idle":"2023-04-19T16:19:08.625157Z","shell.execute_reply.started":"2023-04-19T16:19:05.174558Z","shell.execute_reply":"2023-04-19T16:19:08.624156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Finally, Kaggle expects a runlength-encoded submission.csv file, so let's output that.","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\ndef rle(output):\n    pixels = np.where(output.flatten() > THRESHOLD, 1, 0).astype(np.uint8)\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] = runs[1::2] - runs[:-1:2]\n    return ' '.join(str(x) for x in runs)\nrle_output = rle(output)\n# This doesn't make too much sense, but let's just output in the required format\n# so notebook works as a submission. :-)\nprint(\"Id,Predicted\\na,\" + rle_output + \"\\nb,\" + rle_output, file=open('submission.csv', 'w'))","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:24:55.928766Z","iopub.execute_input":"2023-04-19T16:24:55.929211Z","iopub.status.idle":"2023-04-19T16:24:56.508569Z","shell.execute_reply.started":"2023-04-19T16:24:55.929174Z","shell.execute_reply":"2023-04-19T16:24:56.507365Z"},"trusted":true},"execution_count":null,"outputs":[]}]}