2_-_Extract_growth_line_pixels.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f9a94e53",
   "metadata": {},
   "source": [
    "## Extract pixels along (a given) growth axis line\n",
    "\n",
    "Extracting the ring measurements is a two-stage process, and, in this notebook,\n",
    "the first stage is performed. To summarise, the pixel intensity values are\n",
    "extracted along the axis of maximum growth from the segmented patches (the\n",
    "predictions returned from the model). This gives a vector containing values\n",
    "which correspond to, according to the neural network model, the probability of\n",
    "a pixel belonging to a growth ring. Peakfinding, the second stage, is carried out\n",
    "in the [next notebook](./3_-_Peakfinding_for_ring_widths.ipynb). As with the\n",
    "predictions, growth line pixel extraction is carried out separately for the\n",
    "LOIO and final experiments, although the methodology is identical between the\n",
    "two. A square `k` by `k` patch, centred on a location along the growth axis, is\n",
    "extracted from the model's predictions. The maximum value of this patch is\n",
    "taken as the value assigned to the growth axis at the corresponding location.\n",
    "This is repeated for all locations along the growth axis, and results in a\n",
    "vector of pixel intensities (pseudo probability values of a pixel containing a\n",
    "ring). Further details about this notebook, the files it reads and saves, as\n",
    "well as all other aspects of the project can be found in the\n",
    "[README](./README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0972f1ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import tqdm.notebook\n",
    "import yaml\n",
    "\n",
    "import numpy as np\n",
    "from shellai import preprocessing, tf_util\n",
    "\n",
    "# load the config file\n",
    "with open(\"config.yaml\", \"r\") as fd:\n",
    "    cfg = yaml.safe_load(fd)\n",
    "print('Config file loaded')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfcd97ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # #  Directory settings\n",
    "# base directory -- where we're store things to\n",
    "dir_base = cfg['paths']['local']['base']\n",
    "\n",
    "# directory containing the saved models and their corresponding raw predictions\n",
    "dir_saved_model = cfg['paths']['local']['model']\n",
    "\n",
    "# directory where we wish to save the extracted rings to\n",
    "dir_saved_predictions = os.path.join(dir_base, \"output\")\n",
    "\n",
    "# # # Pixel extraction settings\n",
    "# Note that the pixel extraction works by, for each coordinate along the drawn\n",
    "# mask, we take the mean/max value of a (N, N) sized patch. Empirically, we\n",
    "# found the following values to work well\n",
    "\n",
    "# method: \"mean\" or \"max\", using the mean or maximum value of the patch\n",
    "ringfinding_method = \"max\"\n",
    "\n",
    "# size of the patch from which to extract the ring\n",
    "ringfinding_patch_size = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75abfefc",
   "metadata": {},
   "source": [
    "### Perform the extraction for the leave-one-image-out (LOIO) models\n",
    "Note that we have to do this separately for both the LOIO models and the final model due to the different file and directory structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7df45730",
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings for the LOIO experiments\n",
    "\n",
    "# location of the real images\n",
    "base_image_dir = os.path.join(dir_base, cfg['images']['train_dir'])\n",
    "\n",
    "# mask of what we want the predictions to be called\n",
    "save_name_mask = \"loo_{image_name}_all_predictions.npz\"\n",
    "\n",
    "# list of the images used in training. each one corresponds to a model where\n",
    "# the model was trained on all but the named image.\n",
    "image_names = cfg['images']['train_images']\n",
    "\n",
    "# mask of the model predictions\n",
    "pred_file_mask = 'model_checkpoint_*_predictions.npz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1947b2c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for image_name in tqdm.notebook.tqdm(image_names):\n",
    "    save_path = os.path.join(\n",
    "        dir_saved_predictions, \n",
    "        save_name_mask.format(image_name=image_name)\n",
    "    )\n",
    "\n",
    "    if os.path.exists(save_path):\n",
    "        print(f'Already exists, skipping: {save_path:s}')\n",
    "        continue\n",
    "\n",
    "    # load the image and ground-truth segmentation\n",
    "    image, _, ring_mask_full = preprocessing.load_image_data(\n",
    "        base_image_dir, image_name\n",
    "    )\n",
    "    ring_mask = preprocessing.threshold_drawn_mask_and_skeletonize(\n",
    "        ring_mask_full, sparse=False\n",
    "    )\n",
    "    # list of each model's predictions (different saved model per number of epochs)\n",
    "    pred_paths = glob.glob(\n",
    "        os.path.join(dir_saved_model, image_name, pred_file_mask)\n",
    "    )\n",
    "\n",
    "    # get the epoch numbers\n",
    "    epochs = [\n",
    "        int(os.path.basename(pred_path).split('_')[2])\n",
    "        for pred_path in pred_paths\n",
    "    ]\n",
    "\n",
    "    # sort the paths by epoch number\n",
    "    sorted_inds = np.argsort(epochs)\n",
    "    epochs = [epochs[i] for i in sorted_inds]\n",
    "    pred_paths = [pred_paths[i] for i in sorted_inds]\n",
    "\n",
    "    # storage for saving\n",
    "    ring_data = {\"preds\": {e: None for e in epochs}, \"gt\": None}\n",
    "\n",
    "    # just to ignore bug\n",
    "    patch_coords = flc = None\n",
    "\n",
    "    for epoch, pred_path in tqdm.notebook.tqdm(\n",
    "        zip(epochs, pred_paths), total=len(epochs), leave=False\n",
    "    ):\n",
    "        # load the predictions, coordinates along the drawn line, and coordinates\n",
    "        # of each extracted patch\n",
    "        with np.load(pred_path, allow_pickle=True) as fd:\n",
    "            predictions = fd['predictions']\n",
    "            flc = fd['full_line_coords']\n",
    "            patch_coords = fd['patch_coords'] # [c0, c1, r0, r1]\n",
    "            # patch_centres = fd['patch_centres'] # [idx0, idx1]\n",
    "        \n",
    "        # here, 'gt' refers to pixels extracted using the ground-truth segmentation\n",
    "        if ring_data[\"gt\"] is None:\n",
    "            ring_data[\"gt\"] = tf_util.extract_rings_patchbased(\n",
    "                line_coords=flc,\n",
    "                image=ring_mask, #type: ignore\n",
    "                patch_size=ringfinding_patch_size,\n",
    "                method=ringfinding_method\n",
    "            )\n",
    "\n",
    "        # create the prediction image and store it. this maps the predicted patches\n",
    "        # into an image of a predefined shape\n",
    "        predicted_mask, _ = tf_util.create_patch_image(\n",
    "            patches=predictions,\n",
    "            patch_coords=patch_coords,\n",
    "            image_shape=image.shape\n",
    "        )\n",
    "        \n",
    "        # as with the gt, extract the pixels along given coordinates\n",
    "        ring_data[\"preds\"][epoch] = tf_util.extract_rings_patchbased(\n",
    "            line_coords=flc,\n",
    "            image=predicted_mask,\n",
    "            patch_size=ringfinding_patch_size,\n",
    "            method=ringfinding_method\n",
    "        )\n",
    "\n",
    "    # save everything\n",
    "    print(f'Saving: {save_path:s}')\n",
    "    np.savez(\n",
    "        save_path,\n",
    "        ring_data=ring_data,\n",
    "        image_name=image_name,\n",
    "        epochs=epochs,\n",
    "        flc=flc,\n",
    "        patch_coordintes=patch_coords,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84f405f0",
   "metadata": {},
   "source": [
    "## Repeat for the final model\n",
    "Training of the final model was carried out using the ten training images from the LOIO experiments. Testing is carried out on images with no ground truth segmentation available and, therefore, no ground truth pixel values can be extracted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c710602a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings for the LOIO experiments\n",
    "\n",
    "# location of the real images\n",
    "base_image_dir = os.path.join(dir_base, cfg['images']['test_dir'])\n",
    "\n",
    "# mask of what we want the predictions to be called\n",
    "save_name_mask = \"final_{image_name}_all_predictions.npz\"\n",
    "\n",
    "# list of the images used in training. each one corresponds to a model where\n",
    "# the model was trained on all but the named image.\n",
    "image_names = cfg['images']['test_images']\n",
    "\n",
    "# mask of the model predictions - note that segmentations were carried each set\n",
    "# of model checkpoints, for each test image\n",
    "pred_file_mask = 'model_checkpoint_*_predictions_{image_name:s}.npz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42e10467",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_name = image_names[0]\n",
    "save_path = os.path.join(\n",
    "    dir_saved_predictions, \n",
    "    save_name_mask.format(image_name=image_name)\n",
    ")\n",
    "print(save_path)\n",
    "\n",
    "with np.load(save_path, allow_pickle=True) as fd:\n",
    "    print(fd.files)\n",
    "    ring_data = fd['ring_data'].item()\n",
    "    print(fd['image_name'])\n",
    "    print(fd['epochs'])\n",
    "    flc = fd['flc']\n",
    "    patch_coordintes = fd['patch_coordintes']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54df348f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for image_name in tqdm.notebook.tqdm(image_names):\n",
    "    save_path = os.path.join(\n",
    "        dir_saved_predictions, \n",
    "        save_name_mask.format(image_name=image_name)\n",
    "    )\n",
    "\n",
    "    if os.path.exists(save_path):\n",
    "        print(f'Already exists, skipping: {save_path:s}')\n",
    "        continue\n",
    "\n",
    "    # load the image and ground-truth segmentation\n",
    "    image, _ = preprocessing.load_image_data(\n",
    "        base_image_dir, image_name, no_rings=True\n",
    "    )\n",
    "\n",
    "    # list of each model's predictions (different saved model per number of epochs)\n",
    "    pred_paths = glob.glob(\n",
    "        os.path.join(\n",
    "            dir_saved_model, \n",
    "            \"final_model\", \n",
    "            pred_file_mask.format(image_name=image_name)\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    # get the epoch numbers\n",
    "    epochs = [\n",
    "        int(os.path.basename(pred_path).split('_')[2])\n",
    "        for pred_path in pred_paths\n",
    "    ]\n",
    "\n",
    "    # sort the paths by epoch number\n",
    "    sorted_inds = np.argsort(epochs)\n",
    "    epochs = [epochs[i] for i in sorted_inds]\n",
    "    pred_paths = [pred_paths[i] for i in sorted_inds]\n",
    "\n",
    "    # storage for saving\n",
    "    ring_data = {\"preds\": {e: None for e in epochs}, \"gt\": None}\n",
    "\n",
    "    # just to ignore bug\n",
    "    patch_coords = flc = None\n",
    "\n",
    "    for epoch, pred_path in tqdm.notebook.tqdm(\n",
    "        zip(epochs, pred_paths), total=len(epochs), leave=False\n",
    "    ):\n",
    "        # load the predictions, coordinates along the drawn line, and coordinates\n",
    "        # of each extracted patch\n",
    "        with np.load(pred_path, allow_pickle=True) as fd:\n",
    "            predictions = fd['predictions']\n",
    "            flc = fd['full_line_coords']\n",
    "            patch_coords = fd['patch_coords'] # [c0, c1, r0, r1]\n",
    "            # patch_centres = fd['patch_centres'] # [idx0, idx1]\n",
    "\n",
    "        # create the prediction image and store it. this maps the predicted \n",
    "        # patches into an image of a predefined shape\n",
    "        predicted_mask, _ = tf_util.create_patch_image(\n",
    "            patches=predictions,\n",
    "            patch_coords=patch_coords,\n",
    "            image_shape=image.shape\n",
    "        )\n",
    "        \n",
    "        # extract the pixels along given coordinates\n",
    "        ring_data[\"preds\"][epoch] = tf_util.extract_rings_patchbased(\n",
    "            line_coords=flc,\n",
    "            image=predicted_mask,\n",
    "            patch_size=ringfinding_patch_size,\n",
    "            method=ringfinding_method\n",
    "        )\n",
    "\n",
    "    # save everything\n",
    "    print(f'Saving: {save_path:s}')\n",
    "    np.savez(\n",
    "        save_path,\n",
    "        ring_data=ring_data,\n",
    "        image_name=image_name,\n",
    "        epochs=epochs,\n",
    "        flc=flc,\n",
    "        patch_coordintes=patch_coords,\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}