{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Semantic Segmentation of Nucleii\n\n* This is an example of semantic segmentaton using the the Data Science Bowl 2018 dataset; cell nucleii images are labeled using image masks.\n* While the competition was looking for instance segmentation (each nucleii labeled seperately), here we tackle the less complex problem of semantic segmentation.\n* Using the TensorFlow/Keras framework, we train a U-Net neural network to classify the areas of the images that are nucleii, essentially classifying each pixel.","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport multiprocessing\nimport glob\nimport datetime\nimport random\nimport time\nimport numpy as np\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.preprocessing.image import load_img, ImageDataGenerator\nimport PIL\nfrom PIL import Image\nfrom PIL import ImageChops\nfrom PIL import ImageOps\nimport pandas as pd\nfrom IPython.display import FileLink","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 1. Setup Constants and Data, Start TensorBoard","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"TRAIN_HOME_DIR = \"/kaggle/working/stage1_train\"\nTEST_HOME_DIR = \"/kaggle/working/stage1_test\"\nMODEL_CHECKPOINT_NAME = \"/kaggle/working/nucleii_segmentation.h5\"\nTB_LOGDIR = \"/kaggle/working/tensorboardlogs\"\nTRAIN_SPLIT = 0.8\nIMG_SIZE = (512, 512)\nNUM_CLASSES = 2\nBATCH_SIZE = 4\nEPOCHS = 50\nAUGMENTATION_ON = False","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Unzip the training data, download and unzip ngrok to run TensorBoard later\n!mkdir -p /kaggle/working/stage1_train && unzip -n -q /kaggle/input/data-science-bowl-2018/stage1_train.zip -d /kaggle/working/stage1_train\n# !mkdir -p /kaggle/working/stage1_test  && unzip -n -q /kaggle/input/data-science-bowl-2018/stage1_test.zip -d /kaggle/working/stage1_test\n!wget -Nq https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip; unzip -n -q ngrok-stable-linux-amd64.zip -d /kaggle/working","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Start Up TensorBoard and ngrok\n# Start TensorBoard, `ngrok` opens a tunnel to our Kaggle session to connect to TensorBoard\n# A clickable URL appears in the output below\n\npool = multiprocessing.Pool(processes = 10)\nresults_of_processes = [pool.apply_async(os.system, args=(cmd, ), callback = None )\n                        for cmd in [\n                        f\"tensorboard --logdir {TB_LOGDIR} --host 0.0.0.0 --port 6006 &\",\n                        \"./ngrok http 6006 &\"\n                        ]]\ntime.sleep(2)\n!curl -s http://localhost:4040/api/tunnels | python3 -c \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 2. Define Functions","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def generate_combined_mask(home_dir):\n    path_list = [d.path for d in os.scandir(home_dir) if d.is_dir()]\n    for img_path in tqdm(iterable=path_list, desc=\"Processing mask files\"):\n        searchpath = os.path.join(img_path, \"masks\", \"*.png\")\n        masklist = glob.glob(searchpath)\n        firstmask = Image.open(masklist[0], 'r')\n        img_w, img_h = firstmask.size\n        background_image = Image.new('L', (img_w, img_h), 0)\n        for m in masklist:\n            background_image = ImageChops.lighter(background_image, Image.open(m))\n        new_mask_dir = os.path.join(img_path, \"masks2\")\n        os.makedirs(new_mask_dir, exist_ok=True)\n        new_mask_path = os.path.join(new_mask_dir, \"newmask.png\")\n        background_image.save(new_mask_path)\n    image_list = [f.path for i in path_list for f in os.scandir(os.path.join(i, \"images\")) if f.is_file()]\n    masks_list = [f.path for i in path_list for f in os.scandir(os.path.join(i, \"masks2\")) if f.is_file()]\n    return image_list, masks_list\n\nclass Nucleii(keras.utils.Sequence):\n    # Helper to turn images into Sequence object for TF model\n    def __init__(self, batch_size, img_size, input_img_paths, target_img_paths, img_dtype=\"float32\", tgt_dtype=\"uint8\"):\n        self.batch_size = batch_size\n        self.img_size = img_size\n        self.input_img_paths = input_img_paths\n        self.target_img_paths = target_img_paths\n        self.img_dtype = img_dtype\n        self.tgt_dtype = tgt_dtype\n\n    def __len__(self):\n        return len(self.target_img_paths) // self.batch_size\n\n    def __getitem__(self, idx):\n        # Returns tuple (input, target) correspond to batch #idx.\n        i = idx * self.batch_size\n        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]\n        batch_target_img_paths = self.target_img_paths[i : i + self.batch_size]\n        x = np.zeros((self.batch_size,) + self.img_size + (3,), dtype=self.img_dtype)\n        for j, path in enumerate(batch_input_img_paths):\n            img = load_img(path, target_size=self.img_size)\n            x[j] = img\n        y = np.zeros((self.batch_size,) + self.img_size + (1,), dtype=self.tgt_dtype)\n        for j, path in enumerate(batch_target_img_paths):\n            img = load_img(path, target_size=self.img_size, color_mode=\"grayscale\")\n            tgt_array = np.array(img) / 255\n            y[j] = np.expand_dims(tgt_array, 2)\n        return x, y\n\ndef image_and_mask_generator(image_list, masks_list, generator_args, image_size, batch_size):\n    generator_list = []\n    seed = 1\n    colormode = {0:\"rgb\", 1:\"grayscale\"}\n    for i, j in enumerate([image_list, masks_list]):\n        dtype = (\"uint8\" if i==1 else None) # for mask dtype\n        generator_args[\"dtype\"] = dtype\n        datagen = ImageDataGenerator(**generator_args)\n        generator = datagen.flow_from_dataframe(\n        dataframe=pd.DataFrame(j),\n        directory=None,\n        x_col=0,\n        target_size=image_size,\n        color_mode=colormode[i],\n        class_mode=None,\n        batch_size=batch_size,\n        seed=seed)\n        generator_list.append(generator)\n    return zip(*generator_list)\n    \ndef get_model(img_size, num_classes):\n    inputs = keras.Input(shape=img_size + (3,))\n\n    ### [First half of the network: downsampling inputs] ###\n    # Entry block\n    x = layers.Conv2D(32, 3, strides=2, padding=\"same\")(inputs)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation(\"relu\")(x)\n\n    previous_block_activation = x  # Set aside residual\n\n    # Blocks 1, 2, 3 are identical apart from the feature depth.\n    for filters in [64, 128, 256]:\n        x = layers.Activation(\"relu\")(x)\n        x = layers.SeparableConv2D(filters, 3, padding=\"same\")(x)\n        x = layers.BatchNormalization()(x)\n\n        x = layers.Activation(\"relu\")(x)\n        x = layers.SeparableConv2D(filters, 3, padding=\"same\")(x)\n        x = layers.BatchNormalization()(x)\n\n        x = layers.MaxPooling2D(3, strides=2, padding=\"same\")(x)\n\n        # Project residual\n        residual = layers.Conv2D(filters, 1, strides=2, padding=\"same\")(\n            previous_block_activation\n        )\n        x = layers.add([x, residual])  # Add back residual\n        previous_block_activation = x  # Set aside next residual\n\n    ### [Second half of the network: upsampling inputs] ###\n\n    previous_block_activation = x  # Set aside residual\n\n    for filters in [256, 128, 64, 32]:\n        x = layers.Activation(\"relu\")(x)\n        x = layers.Conv2DTranspose(filters, 3, padding=\"same\")(x)\n        x = layers.BatchNormalization()(x)\n\n        x = layers.Activation(\"relu\")(x)\n        x = layers.Conv2DTranspose(filters, 3, padding=\"same\")(x)\n        x = layers.BatchNormalization()(x)\n\n        x = layers.UpSampling2D(2)(x)\n\n        # Project residual\n        residual = layers.UpSampling2D(2)(previous_block_activation)\n        residual = layers.Conv2D(filters, 1, padding=\"same\")(residual)\n        x = layers.add([x, residual])  # Add back residual\n        previous_block_activation = x  # Set aside next residual\n\n    # Add a per-pixel classification layer\n    outputs = layers.Conv2D(num_classes, 3, activation=\"sigmoid\", padding=\"same\")(x)\n\n    # Define the model\n    model = keras.Model(inputs, outputs)\n    return model\n\n# This is a bug fix for the Keras MeanIoU metric \n# From https://stackoverflow.com/questions/61824470/dimensions-mismatch-error-when-using-tf-metrics-meaniou-with-sparsecategorical\nclass UpdatedMeanIoU(tf.keras.metrics.MeanIoU):\n  def __init__(self,\n               y_true=None,\n               y_pred=None,\n               num_classes=None,\n               name=None,\n               dtype=None):\n    super(UpdatedMeanIoU, self).__init__(num_classes = num_classes,name=name, dtype=dtype)\n\n  def update_state(self, y_true, y_pred, sample_weight=None):\n    y_pred = tf.math.argmax(y_pred, axis=-1)\n    return super().update_state(y_true, y_pred, sample_weight)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3. Prepare Combined Masks\nCreate a single combined mask per image so we can do easier semantic segmentation instead of the more difficult instance segmentation. Return a list of image paths and mask paths","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"image_list, masks_list = generate_combined_mask(TRAIN_HOME_DIR)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4. Create Image Sequence Objects/Generators","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"##--> DATA SETUP: Split our image and mask paths into training/validation sets\nval_samples = int(len(image_list) * (1 - TRAIN_SPLIT))\nrandom.Random(1337).shuffle(image_list)\nrandom.Random(1337).shuffle(masks_list)\ntrain_image_list = image_list[:-val_samples]\ntrain_masks_list = masks_list[:-val_samples]\nval_image_list = image_list[-val_samples:]\nval_masks_list = masks_list[-val_samples:]\n\n# NO AUGMENTATION: Instantiate data Sequence objects for each split\ntrain_seq = Nucleii(BATCH_SIZE, IMG_SIZE, train_image_list, train_masks_list)\nval_seq = Nucleii(BATCH_SIZE, IMG_SIZE, val_image_list, val_masks_list)\n\n# ADD AUGMENTATION: Create generator objects that can create infinite augmented images from base dataset\ntrain_data_gen_args  =  dict(rescale=1./255,\n                        shear_range=0.5,\n                        rotation_range=50,\n                        zoom_range=0.2,\n                        width_shift_range=0.2,\n                        height_shift_range=0.2,\n                        fill_mode='reflect'\n                        )\n                          \nval_data_gen_args = dict(rescale=1./255,\n                        )\n\ntrain_gen_aug = image_and_mask_generator(train_image_list, train_masks_list, train_data_gen_args, IMG_SIZE, BATCH_SIZE)\nval_gen_aug = image_and_mask_generator(val_image_list, val_masks_list, val_data_gen_args, IMG_SIZE, BATCH_SIZE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 5. Create Model, Begin Training\n* Define callbacks (incl. TensorBoard, Early Stopping)\n* Define metrics (here MeanIoU per image)\n* Compile model, specifying optimizer & loss function\n* Begin training for specified number of epochs","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Free up RAM in case the model definition cells were run multiple times\nkeras.backend.clear_session()\n\n# Set up logging directory, callback functions, and metrics (for TensorBoard)\nlog_dir = os.path.join(TB_LOGDIR, datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\ncallbacks = [\n    keras.callbacks.ModelCheckpoint(MODEL_CHECKPOINT_NAME, save_best_only=True),\n    keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=5, write_graph=True, write_images=True, embeddings_freq=5),\n    tf.keras.callbacks.EarlyStopping(monitor=\"val_loss\", verbose=1, patience=7)\n]\nmetrics = [\n    UpdatedMeanIoU(num_classes=NUM_CLASSES), # bug fix for tf.keras.metrics.MeanIoU, see above\n]\n\n# Set up model layers with get_model(), choose optimizer and loss function in compile step\nmodel = get_model(IMG_SIZE, NUM_CLASSES)\n#model.summary()\nmodel.compile(optimizer=\"rmsprop\", loss=\"sparse_categorical_crossentropy\", metrics=metrics)\n\n# Train the model, doing validation at the end of each epoch\nif(AUGMENTATION_ON):\n    # WITH AUGMENTATION\n    print(f\"Beginning training for {EPOCHS} epochs, batch size: {BATCH_SIZE}, augmentation: {AUGMENTATION_ON}\")\n    model.fit(train_gen_aug, steps_per_epoch=1500//BATCH_SIZE, epochs=EPOCHS, validation_data=val_seq, callbacks=callbacks)\nelse:\n    # WITHOUT AUGMENTATION\n    print(f\"Beginning training for {EPOCHS} epochs, batch size: {BATCH_SIZE}, augmentation: {AUGMENTATION_ON}\")\n    model.fit(train_seq, epochs=EPOCHS, validation_data=val_seq, callbacks=callbacks)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 6. Predict and Validate Output\nNow that our model is trained, let's predict against some sample data and compare the ground truth mask with our predicted mask.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"val_preds = model.predict(val_seq)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Randomly select four images from validation data\n# Display input image, input mask, and predicted mask\ndef plot_images(image_list, mask_list, predictions, sample_size):\n    i = np.random.randint(0, len(image_list)-1, size=sample_size)\n    f, axarr = plt.subplots(sample_size//2, 6, figsize=(24,int(sample_size*2)))\n    axarr = axarr.flatten()\n    _ = [a.set_axis_off() for a in axarr.ravel()]\n    for x in range(sample_size):\n        axarr[3*x].imshow(load_img(image_list[i[x]]))\n        axarr[3*x].set_title(\"Original Image\")\n        axarr[3*x+1].imshow(load_img(mask_list[i[x]]), cmap=\"gray\")\n        axarr[3*x+1].set_title(\"Ground Truth Mask\")\n        pred_mask = np.argmax(predictions[i[x]], axis=-1)\n        pred_mask = np.expand_dims(pred_mask, axis=-1)\n        pred_img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(pred_mask))\n        axarr[3*x+2].imshow(pred_img, cmap=\"gray\")\n        axarr[3*x+2].set_title(\"Predicted Mask\")\n\nVAL_IMG_SAMPLE = 6\nplot_images(val_image_list, val_masks_list, val_preds, VAL_IMG_SAMPLE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Utility to download TensorBoard logs to upload to TensorBoard.dev","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm -f /kaggle/working/tensorboardlogs.tar; cd /kaggle/working; tar czf tensorboardlogs.tar.gz tensorboardlogs\n# FileLink(r'tensorboardlogs.tar.gz')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"###### List of Runs\n* `20200802-225855`: 50 epochs, batch size 4, no augmentation, image size 256,256 \n* `20200802-231718`: 50 epochs, batch size 4, no augmentation, image size 256,256\n* `20200802-232838`: 50 epochs, batch size 8, no augmentation, image size 256,256\n* `20200802-233711`: 50 epochs, batch size 2, no augmentation, image size 256,256\n* optimal --> batch size 4\n* `20200803-001835`: 50 epochs, batch size 4, no augmentation, image size 512,512\n* improved --> image size 512,512\n* `20200803-220200`: 50 epochs, batch size 4, no augmentation, image size 512,512, early stopping at epoch 26\n* `20200803-231343`: 50 epochs, batch size 4, no augmentation, image size 512,512, early stopping at epoch 29","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# TESTING CODE\ndef plot_images2(image_list, mask_list, predictions, sample_size):\n    i = np.random.randint(0, len(image_list)-1, size=sample_size)\n    f, axarr = plt.subplots(sample_size//2, 6, figsize=(24,int(sample_size*2)))\n    axarr = axarr.flatten()\n    _ = [a.set_axis_off() for a in axarr.ravel()]\n    for x in range(sample_size):\n        axarr[3*x].imshow((image_list[i[x]]))\n        axarr[3*x].set_title(\"Original Image\")\n        axarr[3*x+1].imshow((mask_list[i[x]][:,:,0]), cmap=\"gray\")\n        axarr[3*x+1].set_title(\"Ground Truth Mask\")\n        pred_mask = np.argmax(predictions[i[x]], axis=-1)\n        pred_mask = np.expand_dims(pred_mask, axis=-1)\n        pred_img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(pred_mask))\n        axarr[3*x+2].imshow(pred_img, cmap=\"gray\")\n        axarr[3*x+2].set_title(\"Predicted Mask\")\n\n# aug_batch = next(train_gen_aug)\n# t_img = aug_batch[0]\n# t_msk = aug_batch[1]\n# plot_images2(t_img, t_msk, model.predict(aug_batch), 6)\n# plot_images2(t_img, t_msk, model.predict(aug_batch), 6)","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}