1_-_Create_training_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e709771d",
   "metadata": {},
   "source": [
    "## Extract training images for model training\n",
    "\n",
    "This notebook processes images for training the convolutional neural networks\n",
    "to segment clam rings. Please see the project [README](./README.md) for initial\n",
    "setup of the training data, its file structure and the config file used in all\n",
    "notebooks, as well as a fuller description of the methodology employed in this\n",
    "notebook. The notebook uses the training data defined by the variable\n",
    "`train_dir` in the `images` section of the configuration file. It also ensures\n",
    "that the files and folders in `train_dir` are structured correctly, i.e., each\n",
    "image is in its own self-named folder along with a ring segmentation image and\n",
    "axis of maximum growth image.\n",
    "\n",
    "Patches are then extracted from each image and its corresponding mask. This is\n",
    "carried out by extracting patches that overlap across the entire image with\n",
    "a given stride. In this work we use 256 by 256 pixel patches, with a stride of\n",
    "64 pixels. Patches that lie wholly outside the convex hull of the hand-drawn\n",
    "rings are discarded, along with those patches that contain more than 5\\%\n",
    "background pixels (defined as being all black). These settings are specified in\n",
    "the configuration file. These patches are extracted and stored as\n",
    "[TFRecord files](https://www.tensorflow.org/tutorials/load_data/tfrecord). This\n",
    "file type allows for TensorFlow to stream the data from the hard disk while it\n",
    "is training a model, thereby avoiding having to wait each time to read from\n",
    "disk.\n",
    "\n",
    "The TFRecord files are stored in a folder called\n",
    "`{train_dir}_patches_{ph}x{pw}_stride_{sh}x{sw}`, where `{train_dir}` is the\n",
    "name of the directory containing the training images, `{ph}` and `{pw}`, and\n",
    "`{sh}` and `{sw}` are the patch and stride height and width as defined in the\n",
    "config file. As default, the folder will be named:\n",
    "`image_train_patches_256x256_stride_64x64`, and, if using the provided training\n",
    "images, will contain roughly 21gb of files. Please copy the contents of this\n",
    "folder to a Google Cloud Storage (GCS) bucket and change the corresponding\n",
    "variables in the configuration file (`project_id` and `bucket_name`) to point\n",
    "to the account and corresponding bucket. Please see the\n",
    "[GCS guides](https://cloud.google.com/storage/docs/creating-buckets) for\n",
    "further information on setting up GCS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "054b0eb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config file loaded\n"
     ]
    }
   ],
   "source": [
    "import gc\n",
    "import os\n",
    "import yaml\n",
    "import tqdm.notebook\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from scipy.spatial.qhull import _Qhull\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": 3,
   "id": "2f897dd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found images: ['gg06007', 'gg110010', 'gg110013', 'gg110026', 'gg110049', 'gg110053', 'gg110055', 'gg140022', 'GY0023', 'gy0030']\n"
     ]
    }
   ],
   "source": [
    "# settings\n",
    "\n",
    "# Here we are expecting the folder below to contain folders, one per image,\n",
    "# where each folder contains (at least) three files:\n",
    "# - fname_image.ext   -- image in some format (jpg or tif)\n",
    "# - fname_mask.ext    -- mask denoting the line along which to measure rings\n",
    "# - fname_extended_rings.ext  -- mask of growth rings\n",
    "base_folder = cfg['images']['train_dir']\n",
    "\n",
    "# expected filenames\n",
    "names_expected = ['extended_rings', 'mask', 'image']\n",
    "\n",
    "# get each image name from the folder\n",
    "folder_names = []\n",
    "\n",
    "for name in os.listdir(base_folder):\n",
    "    path = os.path.join(base_folder, name)\n",
    "    \n",
    "    if os.path.isdir(path):\n",
    "        dirlist = os.listdir(path)\n",
    "\n",
    "        missing = False\n",
    "\n",
    "        for expected_filename in names_expected:\n",
    "            try:\n",
    "                preprocessing.get_matching_strings_from_list(\n",
    "                    dirlist, expected_filename, 1\n",
    "                )\n",
    "            except ValueError:\n",
    "                print(\n",
    "                    f\"Missing file for folder '{name:s}': {expected_filename:s}\"\n",
    "                )\n",
    "                missing = True\n",
    "            \n",
    "        if not missing:\n",
    "            folder_names.append(name)\n",
    "\n",
    "print(f'Found images: {folder_names}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b04b7276",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving to: image_train_patches_256x256_stride_64x64\n"
     ]
    }
   ],
   "source": [
    "# patch settings\n",
    "patch_shape = tuple(cfg['training']['patch_shape'])\n",
    "stride = (cfg['training']['stride'], cfg['training']['stride'])\n",
    "\n",
    "# discard patches with more than this proportion as fully black (background)\n",
    "empty_proportion = cfg['training']['empty_patch_proportion']\n",
    "\n",
    "# TFRecord settings -- samples per record file. this should be bigger than\n",
    "# a single batch so that files can be read while the network is training\n",
    "# the equation below is (as default) 256 * 128 = 384\n",
    "num_tfrecord_samples = int(cfg['training']['batch_size'] * 1.5)\n",
    "\n",
    "# directory to save the TFRecord files to\n",
    "save_directory = ''.join([\n",
    "    base_folder,\n",
    "    f\"_patches_{patch_shape[0]}x{patch_shape[1]}\",\n",
    "    f\"_stride_{stride[0]}x{stride[1]}\"\n",
    "])\n",
    "\n",
    "if not os.path.exists(save_directory):\n",
    "    os.makedirs(save_directory)\n",
    "    print(f\"Created save directory: {save_directory:s}\")\n",
    "else:\n",
    "    print(f\"Saving to: {save_directory:s}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b5a46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for image_name in tqdm.notebook.tqdm(folder_names):\n",
    "\n",
    "    # load the image and masks\n",
    "    image, line_mask, ring_mask = preprocessing.load_image_data(\n",
    "        base_folder, image_name\n",
    "    )\n",
    "\n",
    "    # convert the image masks to binary arrays\n",
    "    line_mask = preprocessing.threshold_drawn_mask_and_skeletonize(\n",
    "        line_mask, sparse=False\n",
    "    )\n",
    "    ring_mask = preprocessing.threshold_drawn_mask_and_skeletonize(\n",
    "        ring_mask, sparse=False\n",
    "    )\n",
    "\n",
    "    assert isinstance(line_mask, np.ndarray)\n",
    "    assert isinstance(ring_mask, np.ndarray)\n",
    "\n",
    "    # calculate the convex hull the mask coords\n",
    "    ring_mask_coords = np.stack(np.where(ring_mask)).T\n",
    "    ring_mask_coords = ring_mask_coords.astype('float')\n",
    "\n",
    "    hull = _Qhull(\n",
    "        b\"i\", # type:ignore\n",
    "        points=ring_mask_coords,\n",
    "        options=b\"\",\n",
    "        furthest_site=False,\n",
    "        incremental=False,\n",
    "        interior_point=None,\n",
    "    )\n",
    "\n",
    "    # get the indices of each possible patch of the image such that\n",
    "    # Z[i] = [[c0, r0], [c1, r1]], where the bottom left and upper right\n",
    "    # coordinates of the patch are [c0, r0] and [c1, r1], and where\n",
    "    # c = column index and r = row index\n",
    "    patch_inds = preprocessing.extract_patch_indices(\n",
    "        image.shape[:2], patch_shape, stride\n",
    "    )\n",
    "\n",
    "    # remove patches fully outside the convex hull\n",
    "    patch_inside_hull_mask = preprocessing.patch_corner_in_hull(\n",
    "        patch_inds, hull\n",
    "    )\n",
    "    patch_inds = patch_inds[patch_inside_hull_mask]\n",
    "\n",
    "    # remove patches with a large proportion of empty pixels (all black)\n",
    "    non_empty_mask = preprocessing.remove_empty_patches(\n",
    "        image, patch_inds, proportion=empty_proportion\n",
    "    )\n",
    "    patch_inds = patch_inds[non_empty_mask]\n",
    "\n",
    "    # extract the patches\n",
    "    patch_images = preprocessing.extract_patches(image, patch_inds)\n",
    "    patch_lines = preprocessing.extract_patches(line_mask, patch_inds)\n",
    "    patch_rings = preprocessing.extract_patches(ring_mask, patch_inds)\n",
    "\n",
    "    # delete the images from memory\n",
    "    del (\n",
    "        image, \n",
    "        line_mask, \n",
    "        ring_mask, \n",
    "        ring_mask_coords,\n",
    "        non_empty_mask,\n",
    "        \n",
    "    )\n",
    "    gc.collect()\n",
    "\n",
    "    # turn images into TFRecord files. These are designed to be quick to open\n",
    "    # in tensorflow, and can be loaded while the model is taking a step,\n",
    "    # therefore somewhat avoiding (in theory!) bottlenecks due to hardware IO\n",
    "    total_size = patch_images.shape[0]\n",
    "    num_tfrecords = (total_size // num_tfrecord_samples) + 1\n",
    "\n",
    "    for tfrec_num in range(num_tfrecords):\n",
    "        # work out the next start/end image indices to store in the record\n",
    "        start = tfrec_num * num_tfrecord_samples\n",
    "        end = min(start + num_tfrecord_samples, total_size)\n",
    "\n",
    "        samples = range(start, end)\n",
    "\n",
    "        save_filename = (\n",
    "            f\"{image_name}_{tfrec_num:03d}_{len(samples):d}.tfrec\"\n",
    "        )\n",
    "\n",
    "        with tf.io.TFRecordWriter(\n",
    "            os.path.join(save_directory, save_filename),\n",
    "        ) as writer:\n",
    "\n",
    "            # for each image going into the records\n",
    "            for i in tqdm.notebook.tqdm(samples, leave=False):\n",
    "\n",
    "                # turn it into a record\n",
    "                feature = tf_util.create_tfrecord_feature(\n",
    "                    patch_images[i],\n",
    "                    patch_rings[i],\n",
    "                    patch_lines[i],\n",
    "                    patch_inds[i],\n",
    "                    image_name,\n",
    "                )\n",
    "\n",
    "                # turn it into an 'example'\n",
    "                example = tf.train.Example(\n",
    "                    features=tf.train.Features(feature=feature)\n",
    "                )\n",
    "                \n",
    "                # write to the record\n",
    "                writer.write(example.SerializeToString()) # type:ignore\n",
    "\n",
    "        print(f\"Saved tf records: {save_filename:s}\")"
   ]
  }
 ],
 "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
}