{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1440x360 with 4 Axes>",
      "image/png": "[encoded data, 328496 characters]\n"
     },
     "metadata": {
      "needs_background": "dark"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from scipy.misc import ascent\n",
    "\n",
    "dark_bg = True # Used to adjust matplotlib colors for dark backgrounds\n",
    "# <editor-fold desc=\"Matplotlib config\" defaultstate=\"collapsed\">\n",
    "plt.rcParams[\"image.cmap\"] = \"gray\"\n",
    "if dark_bg:\n",
    "    plt.rcParams[\"text.color\"] = \"white\"\n",
    "plt.rcParams[\"xtick.color\"] = \"gray\"\n",
    "plt.rcParams[\"ytick.color\"] = \"gray\"\n",
    "# </editor-fold>\n",
    "\n",
    "# <editor-fold desc=\"Load images\" defaultstate=\"collapsed\">\n",
    "def gray_scale(im):\n",
    "    return np.dot(im, [0.07, 0.72, 0.21]) # grayscale weighted by human color perception\n",
    "\n",
    "images = [ { \"image\": ascent(), \"name\": \"ascent\" } ]\n",
    "\n",
    "for image_path in os.listdir(\"img\"):\n",
    "    image = Image.open(os.path.join(\"img\", image_path), \"r\")\n",
    "    if len(np.shape(image)) == 3:\n",
    "        image = gray_scale(image)\n",
    "\n",
    "    images.append({ \"image\": image, \"name\": os.path.splitext(os.path.basename(image_path))[0] })\n",
    "\n",
    "_, ax = plt.subplots(ncols=len(images), figsize=(20, 5))\n",
    "# </editor-fold>\n",
    "\n",
    "for i, im in enumerate(images):\n",
    "    ax[i].set_title(im[\"name\"])\n",
    "    ax[i].imshow(im[\"image\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Apply Fourier Transformations"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1440x1440 with 16 Axes>",
      "image/png": "[encoded data, 1097276 characters]\n"
     },
     "metadata": {
      "needs_background": "dark"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from numpy.fft import fft2, fftshift, ifft2\n",
    "from os import makedirs\n",
    "\n",
    "_, ax = plt.subplots(ncols=4, nrows=len(images), figsize=(20, 5 * len(images)))\n",
    "\n",
    "def log_spectrum(im): # Calculates the logarithmic spectrum from Fourier Space\n",
    "    return np.log10(1 + np.abs(im) ** 2)\n",
    "\n",
    "def ilog_spectrum(im, im_angle): # Reverts the logarithmic spectrum back to FS\n",
    "    im_abs = np.sqrt(10 ** im - 1)\n",
    "    return im_abs * np.cos(im_angle) + im_abs * np.sin(im_angle) * 1j\n",
    "\n",
    "for i, im in enumerate(images):\n",
    "    im[\"fourier\"] = (im_fourier := fftshift(fft2(im[\"image\"])))\n",
    "    im[\"fourier_angle\"] = np.angle(im_fourier)\n",
    "    im[\"fourier_log_spec\"] = (im_log_spec := log_spectrum(im_fourier))\n",
    "\n",
    "    ax[i][0].set_title(f\"Original image \\\"{im['name']}\\\"\")\n",
    "    ax[i][0].imshow(im[\"image\"], vmin=0, vmax=255)\n",
    "    ax[i][1].set_title(\"Angle part of FT\")\n",
    "    ax[i][1].imshow(np.angle(im_fourier))\n",
    "    ax[i][2].set_title(\"Logarithmic spectrum of FT\")\n",
    "    ax[i][2].imshow(im_log_spec, vmin=0, vmax=20)\n",
    "    ax[i][3].set_title(\"FT reverted\")\n",
    "    ax[i][3].imshow(abs(ifft2(ilog_spectrum(im_log_spec, im[\"fourier_angle\"]))), vmin=0, vmax=255)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Export logarithmic spectrums"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "export_scale = 20.0 # Should be higher than the spectrum max\n",
    "for im in images:\n",
    "    makedirs(\"spectrum\", exist_ok=True)\n",
    "    plt.imsave(f\"spectrum/{im['name']}.png\", im[\"fourier_log_spec\"] / export_scale * 255.0, vmin=0, vmax=255)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Reimport logarithmic spectrums and display"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1440x1440 with 12 Axes>",
      "image/png": "[encoded data, 1126712 characters]\n"
     },
     "metadata": {
      "needs_background": "dark"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "export_scale = 20.0\n",
    "_, ax = plt.subplots(ncols=3, nrows=len(images), figsize=(20, 5*len(images)))\n",
    "\n",
    "for i, im in enumerate(images):\n",
    "    reimport_spec_im = plt.imread(f\"spectrum/{im['name']}.png\")\n",
    "    reimport_spec = np.zeros((im[\"image\"].shape[0], im[\"image\"].shape[1])) # required because of different color channels\n",
    "    if len(reimport_spec_im.shape) == 3: # if image is written with multiple color channels, we need to reduce it to one\n",
    "        reimport_spec[:,:] = reimport_spec_im[:,:,0]\n",
    "\n",
    "    reimport_spec = reimport_spec * export_scale # image is already loaded normalized to one so no need to divide through 255\n",
    "    reimport_fourier = ilog_spectrum(reimport_spec, im[\"fourier_angle\"])\n",
    "    ax[i][0].set_title(\"Loaded FT spectrum\")\n",
    "    ax[i][0].imshow(reimport_spec, vmin=0, vmax=20)\n",
    "    ax[i][1].set_title(f\"Original image \\\"{im['name']}\\\"\")\n",
    "    ax[i][1].imshow(im[\"image\"], vmin=0, vmax=255)\n",
    "    ax[i][2].set_title(\"Modified and reimported image\")\n",
    "    ax[i][2].imshow(abs(ifft2(reimport_fourier)), vmin=0, vmax=255)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}