{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import packages and find all the *.nid files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q:\\Marketing\\Projects\\Jupyter Notebooks\\Sharpen Image with 1 *.nid files\n",
      "['grid1.nid']\n"
     ]
    }
   ],
   "source": [
    "from NSFopen.read import read as afmreader\n",
    "\n",
    "import os\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "\n",
    "import numpy as np\n",
    "import math\n",
    "import collections\n",
    "\n",
    "import cv2 as cv\n",
    "from PIL import Image\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from mpl_toolkits.axes_grid1.colorbar import colorbar\n",
    "%matplotlib inline\n",
    "\n",
    "font = {'size': 12}\n",
    "matplotlib.rc('font', **font)\n",
    "\n",
    "dirpath = os.getcwd()\n",
    "onlyfiles = [f for f in listdir(dirpath) if isfile(join(dirpath, f)) and f.split(\".\")[-1] == \"nid\"]\n",
    "print(dirpath, \"with\", len(onlyfiles), \"*.nid files\")\n",
    "print(onlyfiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fittingFunc(matrix, xaxis, order):\n",
    "    newmatrix = np.zeros(np.shape(matrix))\n",
    "    for i in range(len(matrix)):\n",
    "        p = np.polyfit(xaxis, matrix[i,:], order)\n",
    "        f = np.poly1d(p)\n",
    "        fit_array = f(xaxis)\n",
    "        newmatrix[i, :] = matrix[i, :] - fit_array \n",
    "    return newmatrix\n",
    "\n",
    "def unsharp_mask(image, kernel_size=5, sigma=1.0, amount=1.0, threshold=0):\n",
    "    if kernel_size%2 != 1:\n",
    "        kernel_size += 1  \n",
    "    blurred = cv.GaussianBlur(image, (kernel_size, kernel_size), sigma)\n",
    "    sharpened = float(amount + 1) * image - float(amount) * blurred\n",
    "    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))\n",
    "    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))\n",
    "    sharpened = sharpened.round().astype(np.uint8)\n",
    "    if threshold > 0:\n",
    "        low_contrast_mask = np.absolute(image - blurred) < threshold\n",
    "        np.copyto(sharpened, image, where=low_contrast_mask)\n",
    "    return sharpened\n",
    "\n",
    "def laplacian_sharp(image, ddepth = cv.CV_16S, kernel_size=3, amount = 0.5):\n",
    "    edges = cv.Laplacian(image, ddepth, ksize=kernel_size)\n",
    "    sharpened = (1-amount)*image - amount*edges\n",
    "    armax = np.max(sharpened)\n",
    "    armin = np.min(sharpened)\n",
    "    sharpened = (sharpened - armin)/(armax-armin)*255\n",
    "    return sharpened"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define what to plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "filename = 'grid1.nid'\n",
    "signal_to_plot = 'Z-Axis' # 'Phase', 'Amplitude', 'Z-Axis', 'Z-Axis Sensor' there could be more, i.e. 2nd lock-in signal, etc.\n",
    "scan_direction = 'Forward' # 'Forward', 'Backward', '2nd Forward', '2nd Backward'\n",
    "what_to_plot = 'Image' # 'Image', 'Spec' or 'Sweep'\n",
    "fitting_order = 1 # order of the fitting polynom\n",
    "fitting = True # Fit or not"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build X, Y axes and read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "alldata = afmreader(filename, verbose = False).data\n",
    "rawdata0 = alldata[what_to_plot][scan_direction][signal_to_plot]*1e9 # nm\n",
    "rawdata = np.array(rawdata0, dtype = 'float')\n",
    "\n",
    "param = afmreader(filename, verbose = False).param\n",
    "xrange = param['X']['range'][0]\n",
    "yrange = param['Y']['range'][0]\n",
    "\n",
    "dict_key = filename.split(\".nid\")[0]\n",
    "\n",
    "xaxis = np.linspace(0, xrange*1E6, len(rawdata[0]))\n",
    "\n",
    "\n",
    "yaxis = np.linspace(0, yrange*1E6, len(rawdata))\n",
    "\n",
    "if fitting:\n",
    "    data = fittingFunc(rawdata, xaxis, fitting_order)\n",
    "else:\n",
    "    data = rawdata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare the figure and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 140192 characters]\n",
      "text/plain": [
       "<Figure size 540x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "force_limits = True # in case of glitches and spikes in the plots, make it False in the beginning\n",
    "\n",
    "xlbl = 'X (um)'\n",
    "ylbl = 'Y (um)'\n",
    "fig, axarr = plt.subplots(1, 1, figsize=(7.5,7))\n",
    "\n",
    "vmin = -40\n",
    "vmax = 40\n",
    "cmap = plt.get_cmap('copper')\n",
    "\n",
    "im = axarr.pcolormesh(xaxis, yaxis, data, cmap = cmap, shading='auto')\n",
    "if force_limits:\n",
    "    im.set_clim(vmin = vmin, vmax = vmax)\n",
    "\n",
    "axarr.set_title(filename.split(\".nid\")[0])\n",
    "axarr.set(xlabel=xlbl, ylabel=ylbl)\n",
    "ax_divider = make_axes_locatable(axarr)\n",
    "cax = ax_divider.append_axes(\"right\", size=\"5%\", pad=\"2%\")\n",
    "p = colorbar(im, cax=cax)\n",
    "p.ax.set_ylabel(\"Z (nm)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Change data to 0 - 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(data)):\n",
    "    for j in range(len(data[0])):\n",
    "        if data[i,j] < vmin:\n",
    "            data[i,j] = vmin\n",
    "        if data[i,j] > vmax:\n",
    "            data[i,j] = vmax\n",
    "            \n",
    "spread = 1/(vmax - vmin)\n",
    "imdata = (data-vmin)*spread*255\n",
    "imdata_new = np.uint8(imdata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sharpen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "img = Image.fromarray(imdata_new, 'L')\n",
    "img.save('img.jpg')\n",
    "\n",
    "img = cv.imread('img.jpg', cv.IMREAD_GRAYSCALE) # Save and read img file, because cv2.imdecode is producing an error\n",
    "\n",
    "img_sharp_g = unsharp_mask(img, kernel_size=7, sigma=1.0, amount=25.0, threshold=0.5)\n",
    "img_sharp_l = laplacian_sharp(img, ddepth = cv.CV_16S, kernel_size=3, amount = 0.6)\n",
    "\n",
    "img_sharp_g = img_sharp_g/spread/255 + vmin # Bring back to the physical range\n",
    "img_sharp_l = img_sharp_l/spread/255 + vmin # Bring back to the physical range"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot original and sharpened together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 370644 characters]\n",
      "text/plain": [
       "<Figure size 1008x360 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "im_im = collections.OrderedDict()\n",
    "axnum = 0\n",
    "fig, axarr = plt.subplots(1, 3, figsize=(14,5))\n",
    "fig.tight_layout(pad=5.0)\n",
    "\n",
    "for ax in axarr:\n",
    "    index = str(axnum)\n",
    "    if axnum == 0:\n",
    "        im_im[str(axnum)] = ax.pcolormesh(xaxis, yaxis, data, cmap = cmap, shading = 'auto')\n",
    "        ax.set_title(filename.split(\".nid\")[0])\n",
    "        ax.set(xlabel=xlbl, ylabel=ylbl)\n",
    "        im_im[str(axnum)].set_clim(vmin = vmin, vmax = vmax)\n",
    "    elif axnum == 1:\n",
    "        im_im[str(axnum)] = ax.pcolormesh(xaxis, yaxis, img_sharp_g, cmap = cmap, shading = 'auto')\n",
    "        ax.set_title(filename.split(\".nid\")[0] + \" unsharp\")\n",
    "        ax.set(xlabel=xlbl, ylabel='')\n",
    "        im_im[str(axnum)].set_clim(vmin = vmin, vmax = vmax)\n",
    "    else:\n",
    "        im_im[str(axnum)] = ax.pcolormesh(xaxis, yaxis, img_sharp_l, cmap = cmap, shading = 'auto')\n",
    "        ax.set_title(filename.split(\".nid\")[0] + \" laplacian\")\n",
    "        ax.set(xlabel=xlbl, ylabel='')\n",
    "        im_im[str(axnum)].set_clim(vmin = vmin/3, vmax = vmax/3)\n",
    "    \n",
    "    ax_divider = make_axes_locatable(ax)\n",
    "    cax = ax_divider.append_axes(\"right\", size=\"5%\", pad=\"2%\")\n",
    "    p = colorbar(im_im[str(axnum)], cax=cax)\n",
    "    if axnum == 2:\n",
    "        p.ax.set_ylabel(\"Z (nm)\")\n",
    "    axnum += 1"
   ]
  }
 ],
 "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
