training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, unicode_literals, absolute_import, division\n",
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "matplotlib.rcParams[\"image.interpolation\"] = 'none'\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "from tifffile import imread\n",
    "from csbdeep.utils import Path, normalize\n",
    "\n",
    "from splinedist import fill_label_holes, random_label_cmap, calculate_extents, gputools_available\n",
    "from splinedist.matching import matching, matching_dataset\n",
    "from splinedist.models import Config2D, SplineDist2D, SplineDistData2D\n",
    "\n",
    "np.random.seed(42)\n",
    "lbl_cmap = random_label_cmap()\n",
    "\n",
    "import splinedist.splinegenerator as sg\n",
    "from splinedist.utils import phi_generator, grid_generator, get_contoursize_max\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data\n",
    "\n",
    "We assume that data has already been downloaded via notebook [1_data.ipynb](1_data.ipynb).  \n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "Training data (for input `X` with associated label masks `Y`) can be provided via lists of numpy arrays, where each image can have a different size. Alternatively, a single numpy array can also be used if all images have the same size. Label images need to be integer-valued.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = sorted(glob('data/dsb2018/train/images/*.tif'))\n",
    "Y = sorted(glob('data/dsb2018/train/masks/*.tif'))\n",
    "assert all(Path(x).name==Path(y).name for x,y in zip(X,Y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = list(map(imread,X))\n",
    "Y = list(map(imread,Y))\n",
    "n_channel = 1 if X[0].ndim == 2 else X[0].shape[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normalize images and fill small label holes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axis_norm = (0,1)   # normalize channels independently\n",
    "# axis_norm = (0,1,2) # normalize channels jointly\n",
    "if n_channel > 1:\n",
    "    print(\"Normalizing image channels %s.\" % ('jointly' if axis_norm is None or 2 in axis_norm else 'independently'))\n",
    "    sys.stdout.flush()\n",
    "\n",
    "X = [normalize(x,1,99.8,axis=axis_norm) for x in tqdm(X)]\n",
    "Y = [fill_label_holes(y) for y in tqdm(Y)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split into train and validation datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(X) > 1, \"not enough training data\"\n",
    "rng = np.random.RandomState(42)\n",
    "ind = rng.permutation(len(X))\n",
    "n_val = max(1, int(round(0.15 * len(ind))))\n",
    "ind_train, ind_val = ind[:-n_val], ind[-n_val:]\n",
    "X_val, Y_val = [X[i] for i in ind_val]  , [Y[i] for i in ind_val]\n",
    "X_trn, Y_trn = [X[i] for i in ind_train], [Y[i] for i in ind_train] \n",
    "print('number of images: %3d' % len(X))\n",
    "print('- training:       %3d' % len(X_trn))\n",
    "print('- validation:     %3d' % len(X_val))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training data consists of pairs of input image and label instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_img_label(img, lbl, img_title=\"image\", lbl_title=\"label\", **kwargs):\n",
    "    fig, (ai,al) = plt.subplots(1,2, figsize=(12,5), gridspec_kw=dict(width_ratios=(1.25,1)))\n",
    "    im = ai.imshow(img, cmap='gray', clim=(0,1))\n",
    "    ai.set_title(img_title)    \n",
    "    fig.colorbar(im, ax=ai)\n",
    "    al.imshow(lbl, cmap=lbl_cmap)\n",
    "    al.set_title(lbl_title)\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = min(9, len(X)-1)\n",
    "img, lbl = X[i], Y[i]\n",
    "assert img.ndim in (2,3)\n",
    "img = img if (img.ndim==2 or img.shape[-1]==3) else img[...,0]\n",
    "plot_img_label(img,lbl)\n",
    "None;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration\n",
    "\n",
    "A `SplineDist2D` model is specified via a `Config2D` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(Config2D.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# choose the number of control points (M)\n",
    "M = 8\n",
    "n_params = 2 * M\n",
    "\n",
    "# Predict on subsampled grid for increased efficiency and larger field of view\n",
    "grid = (2,2)\n",
    "\n",
    "# compute the size of the largest contour present in the image-set\n",
    "contoursize_max = get_contoursize_max(Y_trn)\n",
    "\n",
    "conf = Config2D (\n",
    "    n_params        = n_params,\n",
    "    grid            = grid,\n",
    "    n_channel_in    = n_channel,\n",
    "    contoursize_max = contoursize_max,\n",
    ")\n",
    "print(conf)\n",
    "vars(conf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'splinedist'\n",
    "model_basedir = 'models'\n",
    "\n",
    "model = SplineDist2D(conf, name=model_name, basedir=model_basedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = model_basedir + '/' + model_name\n",
    "\n",
    "phi_generator(M, conf.contoursize_max, model_path)\n",
    "grid_generator(M, conf.train_patch_size, conf.grid, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check if the neural network has a large enough field of view to see up to the boundary of most objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# median_size = calculate_extents(list(Y), np.median)\n",
    "# fov = np.array(model._axes_tile_overlap('YX'))\n",
    "# print(f\"median object size:      {median_size}\")\n",
    "# print(f\"network field of view :  {fov}\")\n",
    "# if any(median_size > fov):\n",
    "#     print(\"WARNING: median object size larger than field of view of the neural network.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Augmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can define a function/callable that applies augmentation to each batch of the data generator.  \n",
    "We here use an `augmenter` that applies random rotations, flips, and intensity changes, which are typically sensible for (2D) microscopy images (but you can disable augmentation by setting `augmenter = None`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_fliprot(img, mask): \n",
    "    assert img.ndim >= mask.ndim\n",
    "    axes = tuple(range(mask.ndim))\n",
    "    perm = tuple(np.random.permutation(axes))\n",
    "    img = img.transpose(perm + tuple(range(mask.ndim, img.ndim))) \n",
    "    mask = mask.transpose(perm) \n",
    "    for ax in axes: \n",
    "        if np.random.rand() > 0.5:\n",
    "            img = np.flip(img, axis=ax)\n",
    "            mask = np.flip(mask, axis=ax)\n",
    "    return img, mask \n",
    "\n",
    "def random_intensity_change(img):\n",
    "    img = img*np.random.uniform(0.6,2) + np.random.uniform(-0.2,0.2)\n",
    "    return img\n",
    "\n",
    "\n",
    "def augmenter(x, y):\n",
    "    \"\"\"Augmentation of a single input/label image pair.\n",
    "    x is an input image\n",
    "    y is the corresponding ground-truth label image\n",
    "    \"\"\"\n",
    "    x, y = random_fliprot(x, y)\n",
    "    x = random_intensity_change(x)\n",
    "    # add some gaussian noise\n",
    "    sig = 0.02*np.random.uniform(0,1)\n",
    "    x = x + sig*np.random.normal(0,1,x.shape)\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot some augmented examples\n",
    "img, lbl = X[0],Y[0]\n",
    "plot_img_label(img, lbl)\n",
    "for _ in range(3):\n",
    "    img_aug, lbl_aug = augmenter(img,lbl)\n",
    "    plot_img_label(img_aug, lbl_aug, img_title=\"image augmented\", lbl_title=\"label augmented\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.train(X_trn, Y_trn, validation_data=(X_val,Y_val), augmenter=augmenter, epochs = 300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Threshold optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.optimize_thresholds(X_val, Y_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First predict the labels for all validation images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_val_pred = [model.predict_instances(x, n_tiles=model._guess_n_tiles(x), show_tile_progress=False)[0]\n",
    "              for x in tqdm(X_val)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a GT/prediction example "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_img_label(X_val[0],Y_val[0], lbl_title=\"label GT\")\n",
    "plot_img_label(X_val[0],Y_val_pred[0], lbl_title=\"label Pred\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "splinedist",
   "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.8.-1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}