mri-scan-cnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook setup\n",
    "\n",
    "If TensorFlow complains about NUMA node reads from SysFs being -1, try running the following command:\n",
    "\n",
    "`echo 0 | sudo tee -a /sys/bus/pci/devices/0000\\:01\\:00.0/numa_node`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import os\n",
    "import random\n",
    "\n",
    "import ipyplot\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from sklearn.utils import class_weight\n",
    "from tensorflow import keras as tfk\n",
    "from tensorflow.keras import applications as tfka\n",
    "from tensorflow.keras import layers as tfkl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup TensorFlow\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.get_logger().setLevel(\"INFO\")\n",
    "tf.autograph.set_verbosity(0)\n",
    "os.environ[\"TF_FORCE_GPU_ALLOW_GROWTH\"] = \"true\"  # allow gpu memory growth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize variables\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "IMAGES_PATH = \"data/images\"\n",
    "CHECKPOINT_PATH = \"checkpoints\"\n",
    "MODEL_PATH = \"models\"\n",
    "MODEL_NAME = \"model.keras\"\n",
    "VALIDATION_SPLIT = 0.2\n",
    "TEST_SPLIT = 0.1\n",
    "BATCH_SIZE = 8\n",
    "IMAGE_SIZE = (256, 256)\n",
    "INPUT_SHAPE = IMAGE_SIZE + (3,)\n",
    "LEARNING_RATE = 1e-4\n",
    "MIN_LEARNING_RATE = 1e-7\n",
    "TL_EPOCHS = 1\n",
    "FT_EPOCHS = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set seeds\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(SEED)\n",
    "tf.random.set_seed(SEED)\n",
    "random.seed(SEED)\n",
    "os.environ[\"PYTHONHASHSEED\"] = str(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sanity Checks\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert os.path.exists(IMAGES_PATH), f\"Images path {IMAGES_PATH} does not exist\"\n",
    "assert (\n",
    "    0 < VALIDATION_SPLIT + TEST_SPLIT < 1\n",
    "), \"Validation and test split should be less than 1\"\n",
    "assert 0 < VALIDATION_SPLIT < 1, \"Validation split should be less than 1\"\n",
    "assert 0 < TEST_SPLIT < 1, \"Test split should be less than 1\"\n",
    "assert tf.test.is_gpu_available(), \"No GPU available for TF\"\n",
    "assert tf.test.is_built_with_cuda(), \"No CUDA in TF\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create folder structure\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(CHECKPOINT_PATH):\n",
    "    os.makedirs(CHECKPOINT_PATH)\n",
    "    print(f\"Created checkpoint directory {CHECKPOINT_PATH}\")\n",
    "if not os.path.exists(MODEL_PATH):\n",
    "    os.makedirs(MODEL_PATH)\n",
    "    print(f\"Created model directory {MODEL_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define function to clear memory\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clear_memory() -> None:\n",
    "    \"\"\"Function to clear the memory. Defined to be used later\"\"\"\n",
    "    gc.collect()\n",
    "    tfk.backend.clear_session()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data into datasets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_dataset(\n",
    "    ds: tf.data.Dataset,\n",
    "    test_split: float,\n",
    "    validation_split: float,\n",
    ") -> tuple[tf.data.Dataset]:\n",
    "    ds_size = len(ds)\n",
    "    test_size = int(ds_size * test_split)\n",
    "    validation_size = int(ds_size * validation_split)\n",
    "    train_size = ds_size - test_size - validation_size\n",
    "\n",
    "    test_ds = ds.take(test_size)\n",
    "    validation_ds = ds.skip(test_size).take(validation_size)\n",
    "    train_ds = ds.skip(test_size + validation_size).take(train_size)\n",
    "    return train_ds, validation_ds, test_ds\n",
    "\n",
    "\n",
    "ds = tfk.utils.image_dataset_from_directory(\n",
    "    IMAGES_PATH,\n",
    "    labels=\"inferred\",\n",
    "    label_mode=\"categorical\",\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    seed=SEED,\n",
    "    verbose=False,\n",
    ")\n",
    "\n",
    "\n",
    "train_ds, validation_ds, test_ds = split_dataset(ds, TEST_SPLIT, VALIDATION_SPLIT)\n",
    "\n",
    "print(f\"Batch size: {BATCH_SIZE}\")\n",
    "print(\n",
    "    f\"Sizes (in batches): train = {len(train_ds)}, \"\n",
    "    f\"validation = {len(validation_ds)}, test = {len(test_ds)}\"\n",
    ")\n",
    "print(f\"Classes: {ds.class_names}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show some data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract 5 images from each class\n",
    "images_dict = dict()\n",
    "images_to_show = 5\n",
    "for image, label in train_ds:\n",
    "    label_name = ds.class_names[np.argmax(label[0])]\n",
    "\n",
    "    if label_name not in images_dict:\n",
    "        images_dict[label_name] = []\n",
    "\n",
    "    if len(images_dict[label_name]) < images_to_show:\n",
    "        images_dict[label_name].append(image[0].numpy())\n",
    "\n",
    "    if all(len(images) == images_to_show for images in images_dict.values()):\n",
    "        break\n",
    "\n",
    "\n",
    "for label in sorted(images_dict.keys()):\n",
    "    ipyplot.plot_images(\n",
    "        images_dict[label],\n",
    "        labels=[label] * images_to_show,\n",
    "        max_images=images_to_show,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show distribution of classes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_count = {c: 0 for c in ds.class_names}\n",
    "for _, labels in train_ds:\n",
    "    for label in labels:\n",
    "        class_count[ds.class_names[np.argmax(label)]] += 1\n",
    "\n",
    "weights_list = class_weight.compute_class_weight(\n",
    "    class_weight=\"balanced\",\n",
    "    classes=np.unique(ds.class_names),\n",
    "    y=[ds.class_names[np.argmax(label)] for _, labels in train_ds for label in labels],\n",
    ")\n",
    "class_weights = {c: w for c, w in zip(range(len(ds.class_names)), weights_list)}\n",
    "\n",
    "print(\"Number of samples per class in the train dataset:\")\n",
    "for c, n in class_count.items():\n",
    "    print(f\"\\t{c}: {n}\")\n",
    "\n",
    "print(\"Class weights:\")\n",
    "for c_id, w in class_weights.items():\n",
    "    c = ds.class_names[c_id]\n",
    "    print(f\"\\t{c}: {w}\")\n",
    "\n",
    "# plot the distribution of the classes\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.bar(class_count.keys(), class_count.values())\n",
    "plt.xlabel(\"Classes\")\n",
    "plt.ylabel(\"Number of samples\")\n",
    "plt.title(\"Distribution of the classes\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Augmentation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_augmentation_layer = tfk.Sequential(\n",
    "    [\n",
    "        tfkl.Input(shape=INPUT_SHAPE),\n",
    "        tfkl.RandomFlip(\n",
    "            \"horizontal_and_vertical\",\n",
    "            seed=SEED,\n",
    "        ),\n",
    "        tfkl.RandomRotation(\n",
    "            0.5,\n",
    "            fill_mode=\"constant\",\n",
    "            fill_value=0,\n",
    "            seed=SEED,\n",
    "        ),\n",
    "        tfkl.RandomZoom(\n",
    "            0.1,\n",
    "            fill_mode=\"constant\",\n",
    "            fill_value=0,\n",
    "            seed=SEED,\n",
    "        ),\n",
    "        tfkl.RandomTranslation(\n",
    "            0.1,\n",
    "            0.1,\n",
    "            fill_mode=\"constant\",\n",
    "            fill_value=0,\n",
    "            seed=SEED,\n",
    "        ),\n",
    "        tfkl.RandomContrast(\n",
    "            0.2,\n",
    "            seed=SEED,\n",
    "        ),\n",
    "        tfkl.RandomBrightness(\n",
    "            0.2,\n",
    "            seed=SEED,\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# extract some images from the test dataset\n",
    "data = test_ds.take(1)\n",
    "images, labels = next(iter(data))\n",
    "\n",
    "# apply data augmentation\n",
    "augmented_images = data_augmentation_layer(images)[:5]\n",
    "# display the augmented images\n",
    "augmented_image_array = [np.array(image) for image in augmented_images]\n",
    "labels = [ds.class_names[np.argmax(label)] for label in labels]\n",
    "\n",
    "ipyplot.plot_images(augmented_image_array, labels, max_images=5, img_width=150)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Model\n",
    "\n",
    "- The approach I am going for is to use a pretrained model and fine-tune it for this dataset.\n",
    "- [The keras documentation](https://keras.io/api/applications/efficientnet_v2/#efficientnetv2l-function) shows the available models\n",
    "  1. I initially tried the [EfficientNetV2L model](https://keras.io/api/applications/efficientnet_v2/#efficientnetv2l-function) as it has a good score - TOO MUCH MEMORY\n",
    "  2. VGG16?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Definition\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "clear_memory()\n",
    "# load the pretrained model\n",
    "resize_layer = tfkl.Resizing(*IMAGE_SIZE)\n",
    "base_model = tfka.VGG16(\n",
    "    input_shape=(*IMAGE_SIZE, 3),\n",
    "    include_top=False,\n",
    "    classes=len(ds.class_names),\n",
    ")\n",
    "# add preprocessing layer\n",
    "preprocess_input = tfka.vgg16.preprocess_input\n",
    "# freeze the model\n",
    "base_model.trainable = False\n",
    "# create input layer\n",
    "inputs = tfkl.Input(shape=(*IMAGE_SIZE, 3))\n",
    "# create preprocessing layer\n",
    "# create an output network\n",
    "output_layer = tfk.models.Sequential(\n",
    "    [\n",
    "        tfkl.GlobalAveragePooling2D(),\n",
    "        tfkl.Dense(256, activation=\"relu\"),\n",
    "        tfkl.Dropout(0.5),\n",
    "        tfkl.Dense(len(ds.class_names), activation=\"softmax\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "# create the model\n",
    "x = resize_layer(inputs)\n",
    "x = data_augmentation_layer(x)\n",
    "x = preprocess_input(x)\n",
    "x = base_model(x)\n",
    "outputs = output_layer(x)\n",
    "\n",
    "model = tfk.Model(inputs, outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile Model and show summary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile the model\n",
    "model.compile(\n",
    "    optimizer=tfk.optimizers.Adadelta(LEARNING_RATE),\n",
    "    loss=tfk.losses.CategoricalCrossentropy(),\n",
    "    metrics=[tfk.metrics.CategoricalAccuracy()],\n",
    ")\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Callbacks and Checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_filename = os.path.join(CHECKPOINT_PATH, \"checkpoints.keras\")\n",
    "# create the early stopping callback\n",
    "early_stopping = tfk.callbacks.EarlyStopping(\n",
    "    monitor=\"val_categorical_accuracy\",\n",
    "    patience=3,\n",
    "    restore_best_weights=True,\n",
    ")\n",
    "# create the reudce learning rate on plateau callback\n",
    "reduce_lr = tfk.callbacks.ReduceLROnPlateau(\n",
    "    monitor=\"val_categorical_accuracy\",\n",
    "    factor=0.1,\n",
    "    patience=2,\n",
    "    min_lr=MIN_LEARNING_RATE,\n",
    ")\n",
    "# create the model checkpoint callback\n",
    "checkpoint = tfk.callbacks.ModelCheckpoint(\n",
    "    os.path.join(CHECKPOINT_PATH, checkpoint_filename),\n",
    "    monitor=\"val_categorical_accuracy\",\n",
    "    save_best_only=True,\n",
    ")\n",
    "callbacks = [early_stopping, reduce_lr, checkpoint]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define functions to save model and history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model(model: tfk.Model, model_path: str, training_path: str = None) -> None:\n",
    "    \"\"\"Function to save the model and its training history\"\"\"\n",
    "    if not model_path.endswith(\".keras\"):\n",
    "        model_path += \".keras\"\n",
    "    if not training_path:\n",
    "        training_path = model_path.replace(\".keras\", \".npy\")\n",
    "    elif not training_path.endswith(\".npy\"):\n",
    "        training_path += \".npy\"\n",
    "\n",
    "    print(f\"Saving model to {model_path}\")\n",
    "    model.save(model_path)\n",
    "    print(f\"Saving training history to {training_path}\")\n",
    "    np.save(\n",
    "        training_path,\n",
    "        model.history.history,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the top layers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    train_ds,\n",
    "    validation_data=validation_ds,\n",
    "    epochs=TL_EPOCHS,\n",
    "    callbacks=callbacks,\n",
    "    class_weight=class_weights,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine Tune Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unfreeze the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unfreeze the model\n",
    "base_model.trainable = True\n",
    "# re-complie the model\n",
    "model.compile(\n",
    "    optimizer=tfk.optimizers.Adam(LEARNING_RATE),\n",
    "    loss=tfk.losses.CategoricalCrossentropy(),\n",
    "    metrics=[tfk.metrics.CategoricalAccuracy()],\n",
    ")\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    train_ds,\n",
    "    validation_data=validation_ds,\n",
    "    epochs=FT_EPOCHS,\n",
    "    callbacks=callbacks,\n",
    "    class_weight=class_weights,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}