{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## driver notebook for image correction, calculation of elevation maps and stitching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "script that takes raw images in selected path (and subfolders), applies correction, produces Zmaps and stitch them into a final image (.tif). It also produces quality check images and data sheets for the various steps, with volume, mean height and roughness coefficient. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import matplotlib\n",
    "#matplotlib.use('TkAgg')\n",
    "get_ipython().magic('matplotlib inline')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "FlumeAutomationLib='../lib'\n",
    "if not FlumeAutomationLib in sys.path:\n",
    "    sys.path.append(FlumeAutomationLib)\n",
    "    \n",
    "import numpy as np\n",
    "from ipywidgets import widgets,interact,IntProgress\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage import morphology\n",
    "from skimage.morphology import closing, square, reconstruction\n",
    "from skimage import filters\n",
    "from OctCorrection import *\n",
    "from ImageProcessing import *\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib_scalebar.scalebar import ScaleBar\n",
    "from matplotlib import colorbar\n",
    "import os\n",
    "import glob\n",
    "import re\n",
    "import pickle\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "import cv2\n",
    "from tifffile import imsave, imread\n",
    "from numba import jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "srcpath='../sample_scans/';\n",
    "pathexp='../sample_scans/output/';\n",
    "if not os.path.exists(pathexp):\n",
    "     os.makedirs(pathexp)\n",
    "l=glob.glob(srcpath+'*/*.raw')\n",
    "l.sort()\n",
    "l"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "visualize the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size=(900,900,1024)  #define here the size of the images to be processed\n",
    "voxel_size=(10/900,10/900,0.00218) #set here the (x,y,z) dimensions of voxels (here in mm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path=l[0]    #select the image to be displayed\n",
    "ori=np.fromfile(path,dtype=np.uint8).reshape(size)  #load it \n",
    "ori=np.flip(ori, axis=2) #and oriente it correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "crop=500  #find here the best position to crop the image and avoid spurious signals such as the Zspacer\n",
    "bckHeight=100 #find here the position of the line above which the background intensity is calculated\n",
    "\n",
    "fig, ax = plt.subplots(2,figsize=(20,10))\n",
    "\n",
    "sl=1   #select the slice to be displayed (first)\n",
    "im=ori[sl,:,:]\n",
    "ax[0].imshow(im.swapaxes(1,0), cmap='gray')\n",
    "ax[0].set_aspect(2.903/11.111) #set here the aspect ratio, \n",
    "                # you can calculate it from the z spacing and x,y spacing parameters in the .srm files\n",
    "ax[0].axhline(y=bckHeight, color='r', linestyle='-', label='background height')\n",
    "ax[0].axhline(y=crop, color='g', linestyle='-', linewidth=3,label='crop')\n",
    "scalebar = ScaleBar(0.011,'mm') # set here the scale 1 pixel = 11 micron\n",
    "ax[0].axis('off')\n",
    "plt.gca().add_artist(scalebar)\n",
    "ax[0].set_title('slice '+str(sl));\n",
    "\n",
    "\n",
    "sl=899   #select the slice to be displayed (last)\n",
    "im=ori[sl,:,:]\n",
    "ax[1].imshow(im.swapaxes(1,0), cmap='gray')\n",
    "ax[1].set_aspect(2.18/11.111) #set here the aspect ratio, \n",
    "                # you can calculate it from the z spacing and x,y spacing parameters in the .srm files\n",
    "ax[1].axhline(y=bckHeight, color='r', linestyle='-', label='background height')\n",
    "ax[1].axhline(y=crop, color='g', linestyle='-', linewidth=3,label='crop')\n",
    "ax[1].legend(loc=2)\n",
    "scalebar = ScaleBar(0.011,'mm') # set here the scale 1 pixel = 11 micron\n",
    "ax[1].axis('off')\n",
    "plt.gca().add_artist(scalebar)\n",
    "ax[1].set_title('slice '+str(sl));\n",
    "#fig.savefig(pathexp+'original_Ascan.png') #save figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CORRECTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#correct images in the list 'l' and collect them in the dictionary \"CorrImages\"\n",
    "CorrImages={}\n",
    "for path in l:\n",
    "    \n",
    "    ori=np.fromfile(path,dtype=np.uint8).reshape(size)\n",
    "    dat=np.flip(ori, axis=2)\n",
    "    im, pts, selection, Z=Correct(dat, bckHeight=bckHeight, crop=crop,size=dat.shape, \n",
    "                                  retFiltSize=12, GaussFiltSize=55,\n",
    "                                  zHeight=500,modelHeight=400)\n",
    "    CorrImages[path]=im   #append the corrected image in the dictionary \"CorrImages\"\n",
    "    \n",
    "#to save the corrected images, uncimment the two following commands    \n",
    "    #fname=path.split('/')[-1][:-4]  #select string to use as file name for saved corrected images\n",
    "    #imsave(pathexp+fname+'.tif', im.astype('int8'))  #save corrected images as tif files in foled \"pathexp\",\n",
    "                                                        #using \"fname\" as file name\n",
    "    \n",
    "    #produce some quality check plots and images\n",
    "    plotSurf(dat.shape,pts,selection, Z)        #plot selected points and fitted surface \n",
    "    fig=plt.figure(figsize=(3,4), dpi=200)     #plot sample slices of the corrected image (first and last)\n",
    "    plt.subplot(2,1,1)\n",
    "    plt.title(path.split('\\\\')[-1])\n",
    "    plt.imshow(CorrImages[path][1,:,:])\n",
    "    plt.subplot(2,1,2)\n",
    "    plt.imshow(CorrImages[path][-1,:,:])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#visualize a slice of one of the corrected images\n",
    "path=l[0]  #select the image\n",
    "image=CorrImages[path]\n",
    "fig, ax = plt.subplots(2,figsize=(20,10))\n",
    "sl=1\n",
    "ax[0].imshow(image[sl,:,:], cmap='gray')\n",
    "ax[0].set_aspect(voxel_size[2]/voxel_size[0])\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') \n",
    "ax[0].axis('off')\n",
    "ax[0].set_title('slice '+str(sl))\n",
    "\n",
    "sl=899\n",
    "ax[1].imshow(image[sl,:,:], cmap='gray')\n",
    "ax[1].set_aspect(voxel_size[2]/voxel_size[0])\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') \n",
    "ax[1].axis('off')\n",
    "ax[1].set_title('slice '+str(sl))\n",
    "\n",
    "plt.gca().add_artist(scalebar)\n",
    "#fig.savefig(path+'original_Ascan.png') #save figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ELEVATION MAPS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#find best threshold, minimum object size and surface reference for elevation map calculation\n",
    "path=l[0]  #select the image\n",
    "image=CorrImages[path]\n",
    "ref_surf=395  #set here the height of the reference surface from which t calculate the elevation map\n",
    "thr=25     #set here the gray value threshold to calculate the binary mask\n",
    "small_obj=30   #set here the minimum size of objects to be kept in the binary mask\n",
    "\n",
    "fig, ax = plt.subplots(2,figsize=(20,10))\n",
    "sl=1\n",
    "ax[0].imshow(morphology.remove_small_objects(image[sl,:,:]>thr,small_obj), cmap='gray')\n",
    "ax[0].set_aspect(voxel_size[2]/voxel_size[0])\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') \n",
    "ax[0].axis('off')\n",
    "ax[0].set_title('slice '+str(sl))\n",
    "\n",
    "\n",
    "sl=899\n",
    "ax[1].imshow(morphology.remove_small_objects(image[sl,:,:]>thr,small_obj), cmap='gray')\n",
    "ax[1].set_aspect(voxel_size[2]/voxel_size[0])\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') \n",
    "ax[1].axis('off')\n",
    "ax[1].set_title('slice '+str(sl))\n",
    "\n",
    "plt.gca().add_artist(scalebar)\n",
    "ax[1].axhline(y=ref_surf, color='g', linestyle='-', linewidth=3,label='reference surface for elevation map calculation')\n",
    "ax[1].legend(loc=2)\n",
    "#fig.savefig(path+'original_Ascan.png') #save figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate Z maps (Z), volume (V), average height (H) and roughness coefficient (R)\n",
    "cnames=['fname', 'volume','height','roughness']\n",
    "df = pd.DataFrame(columns =cnames )\n",
    "ZMaps={}  #ZMaps are collected in the dictionary ZMaps, with path as key\n",
    "\n",
    "for path in l:\n",
    "    mask,Z,V,H,R=CalcZMap(CorrImages[path][:,:ref_surf],thr,small_obj) \n",
    "    ZMaps[path]=Z\n",
    "    \n",
    "    #uncomment the two following commands to save the elevation maps as .tif files in the destination folder \"pathexp\"\n",
    "    #fname=path.split('/')[-1][:-4]\n",
    "    #imsave(pathexp+fname+'_ZMap.tif',Z.astype('int32'))   #uncomment if you want to directly save the images in pathexp\n",
    "    \n",
    "    df=df.append(pd.DataFrame([[path, V,H,R]], columns = cnames))   #append volume, height and roughness results in a dataframe\n",
    "    print(path)\n",
    "#df\n",
    "#df.to_csv(pathexp+'Volume-Height-Roughness.csv',sep=',')   #uncomment to save the calculated values in a .csv file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#visualize one of the elevation maps\n",
    "image=ZMaps[l[0]]  #select the image to be displayed\n",
    "ticks=[0, 30, 60,90,120]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(20,10))\n",
    "plt.imshow(image, cmap='gray')\n",
    "plt.axis('off')\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') # 1 pixel = 11 micron\n",
    "plt.gca().add_artist(scalebar)\n",
    "cbar=plt.colorbar(ticks=np.multiply(ticks, voxel_size[2]*1000));\n",
    "cbar.ax.set_yticklabels(ticks);\n",
    "cbar.set_label('height (μm)', fontsize=16)\n",
    "#plt.savefig(pathexp+'elev_map.png',dpi=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## STITCHING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "#orientate images to be stitched and put them in a matrix, according to the number of rows and columns\n",
    "n_rows=3\n",
    "n_cols=3\n",
    "images=[np.rot90(ZMaps[path]) for path in ZMaps]\n",
    "M=np.asarray(images).reshape((n_rows,n_cols,images[0].shape[0],images[0].shape[1]))\n",
    "M=np.flip(M,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x33160f28>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 527124 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#check the correct orientation of the images\n",
    "n_rows,n_cols, sy,sx=M.shape\n",
    "b=30\n",
    "Z=np.zeros(((sy+b)*n_rows,(sx+b)*n_cols))\n",
    "for row in range(n_rows):\n",
    "    for col in range(n_cols):\n",
    "        im=M[row,col,:,:]\n",
    "        Z[row*sy+row*b:(1+row)*sy+row*b,col*sx+col*b:(1+col)*sx+col*b]=im\n",
    "fig=plt.figure(figsize=(10,10))\n",
    "plt.imshow(Z, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate coordinates of images in the final stitched image\n",
    "OL=ol_mat(M,0.3)\n",
    "OLv=ol_mat(np.rot90(np.rot90(M, axes=(2,3)), axes=(0,1)),0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 387392 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#stitch and visualize the result\n",
    "Z=stitch(M, OL, OLv,0.3,1)\n",
    "fig=plt.figure(figsize=(10,10))\n",
    "plt.imshow(Z, cmap='gray')\n",
    "plt.axis('off')\n",
    "scalebar = ScaleBar(voxel_size[0],'mm') # 1 pixel = 11 micron\n",
    "plt.gca().add_artist(scalebar)\n",
    "cbar=plt.colorbar(ticks=np.multiply(ticks, voxel_size[2]*1000));\n",
    "cbar.ax.set_yticklabels(ticks);\n",
    "cbar.set_label('height (μm)', fontsize=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#uncomment the two following commands to save the stitched elevation map as .tif files in the destination folder \"pathexp\"\n",
    "    #fname=path.split('/')[-1][:-6]\n",
    "    #imsave(pathexp+fname+'_ZMap.tif',Z.astype('int32'))   #uncomment if you want to directly save the images in pathexp\n",
    "    \n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#stitch with OpenCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images=[np.rot90(ZMaps[path].astype('uint8')) for path in ZMaps]\n",
    "stitcher = cv2.createStitcher(False)\n",
    "result = stitcher.stitch((images))\n",
    "result[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final=cv2.imread('final_image.tif', cv2.IMREAD_GRAYSCALE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
