1st_Tutorial_Wavelets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8ee064a8-a550-4214-aea0-57d20d072ae6",
   "metadata": {
    "id": "ozP33D7YvQ2w"
   },
   "source": [
    "# 1st Tutorial: Wavelets for Multiscale Image Processing\n",
    "\n",
    "This content is largely based on [Gabriel Peyré's Numerical Tours](https://github.com/gpeyre/numerical-tours) 🤗\n",
    "\n",
    "\n",
    "## Building Intuitions\n",
    "\n",
    "In the first part, we get an intuition of what wavelets can do for Image compression, we will compare: \n",
    "\n",
    "* Compressing images by averaging on squares; \n",
    "* Keeping the low-frequency Fourier coefficients; \n",
    "* Keeping the largest Wavelet coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b8e80ec-c50e-4e62-9bb6-ab4ac6366e09",
   "metadata": {
    "id": "DpXWRpYXvesV"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import warnings\n",
    "import pywt # to compute wavelet transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c1588e-d5eb-4da5-9e50-fd35be1386c1",
   "metadata": {
    "id": "hJzNnGCPEGw_"
   },
   "outputs": [],
   "source": [
    "def preprocess(f,n):\n",
    "    \"\"\"\n",
    "    Preprocess a single Image by resizing to a square and normalizing \n",
    "    it in range [0,1]\n",
    "    \n",
    "    ---\n",
    "    Inputs\n",
    "    \n",
    "    - f: PIL Image \n",
    "    - n: image width and height \n",
    "    \n",
    "    Outputs: \n",
    "    - f: Image with values in range [0,1]\n",
    "    \"\"\"\n",
    "    n0 = np.min( f.size )\n",
    "    f = f.crop((0, 0, n0, n0))\n",
    "    f = f.resize((n, n))\n",
    "    f = np.array(f)\n",
    "    f = np.clip(f / np.max(f.flatten()), 0,1)\n",
    "    return f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d1fcbc4-2f86-46f6-a668-f2bbfb8b8f26",
   "metadata": {
    "id": "79aEWKx8vxj_"
   },
   "source": [
    "Load an image from the internet. The image has range of value [0,1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e53cb2b-dde7-4a70-8abd-8b135c4a57ac",
   "metadata": {
    "id": "EcEFsEp_vphQ"
   },
   "outputs": [],
   "source": [
    "from urllib.request import urlopen\n",
    "import io\n",
    "\n",
    "n = 512 # target size\n",
    "\n",
    "# Download an image from wikimedia\n",
    "file_adress = 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c7/Tabby_cat_with_blue_eyes-3336579.jpg/639px-Tabby_cat_with_blue_eyes-3336579.jpg'\n",
    "fd = urlopen(file_adress)\n",
    "\n",
    "# Open and preprocess the image\n",
    "image_file = io.BytesIO(fd.read())\n",
    "f = Image.open(image_file)\n",
    "f = preprocess( f, n )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cab83fd-8ff7-48cf-9677-2d8632a3e03d",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 303
    },
    "id": "grrghnPEwgah",
    "outputId": "491c7140-4c9f-44eb-8e20-787a141f6ee7"
   },
   "outputs": [],
   "source": [
    "# Show the image\n",
    "plt.rcParams['figure.figsize'] = [8, 5]\n",
    "plt.imshow(f)\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8af6d30b-f18e-40c8-94be-fb1bd3432202",
   "metadata": {
    "id": "jp1HJ-E3ZoY2"
   },
   "outputs": [],
   "source": [
    "def approx_pixels(f,r):\n",
    "    \"\"\"\n",
    "    Approximate Image by interpolating pixels\n",
    "    \n",
    "    ---\n",
    "    Inputs: \n",
    "    \n",
    "    - f : PIL Image\n",
    "    - r : compression ratio (smaller than 1) \n",
    "    \n",
    "    Outputs: \n",
    "    \n",
    "    - f : compressed image\n",
    "    \n",
    "    \"\"\"\n",
    "    p = int(n*np.sqrt(r))\n",
    "    # In order to interpolate we need to convert \n",
    "    # the image into unsigned int8: i.e. [0,1] --> [0, 255]  \n",
    "    f = Image.fromarray( np.uint8(f*255) )\n",
    "    f = f.resize((p, p), Image.NEAREST )\n",
    "    f = f.resize((n, n), Image.NEAREST )\n",
    "    # Going back to range [0,1]\n",
    "    f = np.array(f)/255.0\n",
    "    f = np.clip(f,0,1)\n",
    "    \n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd777ef0-7241-48c1-8a18-995cd241319e",
   "metadata": {
    "id": "xHUaUCg_v3WI"
   },
   "outputs": [],
   "source": [
    "def approx_fourier(f,r):\n",
    "    \n",
    "    \"\"\"\n",
    "    Approximate Image by low-frequency Fourier components\n",
    "    \n",
    "    ---\n",
    "    Inputs: \n",
    "    \n",
    "    - f : PIL Image\n",
    "    - r : compression ratio (smaller than 1) \n",
    "    \n",
    "    Outputs: \n",
    "    \n",
    "    - fM : Fourier Compressed image\n",
    "    - shifted_mask : centered Fourier Mask \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    m = int(r*n*n)\n",
    "    # compute a 1D indexing from low to high frequencies\n",
    "    x = np.concatenate( (np.arange(0,int(n/2)), np.arange(-int(n/2),0)) )\n",
    "    D = x[:,None]**2 + x[None,:]**2\n",
    "    # fourier mask , i.e. select low frequency components\n",
    "    mask = D < np.sort(D.flatten())[m]\n",
    "    \n",
    "    # inverse the masked FFT\n",
    "    F = np.fft.fft2(f,axes=(0,1))\n",
    "    #convolve image and mask, in Fourier space it is a product :) \n",
    "    F = F * mask[:,:,None]\n",
    "    \n",
    "    # inverse fourier to get back the image\n",
    "    fM = np.real(np.fft.ifft2(F,axes=(0,1)))\n",
    "    \n",
    "    fM = np.clip(fM,0,1)\n",
    "    # shift mask to the center \n",
    "    shifted_mask = np.fft.fftshift(mask)\n",
    "    \n",
    "    return fM, shifted_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c9058ca-743d-47a2-9784-9e045c578ffe",
   "metadata": {
    "id": "kmwM6geYfZsK"
   },
   "outputs": [],
   "source": [
    "def approx_wavelets(f,r):\n",
    "    \"\"\"\n",
    "    Approximate Image by using wavelets\n",
    "    \n",
    "    ---\n",
    "    Inputs: \n",
    "    \n",
    "    - f : PIL Image\n",
    "    - r : compression ratio (smaller than 1) \n",
    "    \n",
    "    Outputs: \n",
    "    \n",
    "    - f1 : Wavelets Compressed image\n",
    "    - mask : thre\n",
    "    \n",
    "    \"\"\"\n",
    "    m = int(r*n*n)\n",
    "    # Select specific wavelet, see pywt docs for more info\n",
    "    wavetype = 'bior3.5'\n",
    "    nlvel = int( np.floor(np.log2(n))-3 )\n",
    "    f1 = np.zeros( f.shape );\n",
    "    \n",
    "    # iterate through channels \n",
    "    for k in range(0, 3):\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter(\"ignore\") # pywlet issue a warning for too low level decomposition\n",
    "            fW = pywt.wavedec2(f[:,:,k], wavetype, level=nlvel)\n",
    "        # turn coefficient into a convenient array\n",
    "        fWI, coeff_slices = pywt.coeffs_to_array(fW)\n",
    "        D = np.sort( np.abs(fWI.flatten()) ) # sorted coefficients magnitude\n",
    "        fWI = fWI * (abs(fWI)>D[-m-1]) # threshold\n",
    "        fW1 = pywt.array_to_coeffs(fWI,coeff_slices, output_format='wavedec2')\n",
    "        f1[:,:,k] = pywt.waverec2( fW1, wavetype)\n",
    "        \n",
    "    mask = (abs(fWI)>D[-m-1])\n",
    "    f1 = np.clip(np.array(f1),0,1)\n",
    "    \n",
    "    return f1, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2934561f-8388-450c-aaa5-e759514344f1",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 574
    },
    "id": "U3D3JTGowwSy",
    "outputId": "afbd99a2-c2af-4ff9-b5a6-a1e4e2acf81f"
   },
   "outputs": [],
   "source": [
    "def render_display(r):\n",
    "    fP = approx_pixels(f,r)\n",
    "    fF, mF = approx_fourier(f,r)\n",
    "    fW, mW = approx_wavelets(f,r)\n",
    "    \n",
    "    plt.rcParams['figure.figsize'] = [16, 10]\n",
    "    \n",
    "    plt.subplot(2,3,1)\n",
    "    plt.imshow(fP)\n",
    "    plt.axis('off')\n",
    "    plt.subplot(2,3,2)\n",
    "    plt.imshow(fF)\n",
    "    plt.axis('off')\n",
    "    plt.subplot(2,3,3)\n",
    "    plt.imshow(fW)\n",
    "    plt.axis('off')\n",
    "    plt.subplot(2,3,5)\n",
    "    plt.imshow(mF)\n",
    "    plt.axis('off')\n",
    "    plt.subplot(2,3,6)\n",
    "    plt.imshow(mW)\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e7ec55-e52f-41c1-bb47-b2fc5f15e23e",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = .01 # compression ratio\n",
    "render_display(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "706700c9-0158-476a-bdba-d07fb1aa7a95",
   "metadata": {},
   "source": [
    "### Interactive Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e399d6d-5604-4d8e-8f27-64b34d58bafe",
   "metadata": {
    "id": "ibdd12DQp44t"
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141c5050-5082-41f8-b4b6-2d145c074168",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 606,
     "referenced_widgets": [
      "73f4ab9cdeba4ab4bdece6ab93cb682e",
      "3ec15f62f137438097e4c181402dbfc7",
      "f0feec4119694a1f97053fc5e56e7652",
      "c116b0dbda3b476c8b309726eb97beff",
      "4727bfcb42da4c74858c27b1e004f425",
      "88b08c59f71f4dd4a0a3fcee5baee0c1",
      "b7c310606af144f78b61e3871926813c"
     ]
    },
    "id": "fRtybQStsV7k",
    "outputId": "bfc26c0a-268a-49ae-e77e-487fea593416"
   },
   "outputs": [],
   "source": [
    "@interact(r=(0.01/2,0.1,0.01/2))\n",
    "def h(r=.05):\n",
    "    render_display(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48406f3b-e93a-4452-bcbf-b060dee5e6b2",
   "metadata": {},
   "source": [
    "## Another Example - Mondrian Figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99499a4b-a391-46f8-b65c-cbe9e9332143",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mondrian(shape=(256, 256), channels = 3, nx=5, ny=8, seed=4):\n",
    "    \n",
    "    \"\"\" Piecewise-constant image (reminiscent of Dutch painter Piet Mondrian's\n",
    "    geometrical period).\n",
    "    \"\"\"\n",
    "    img_final = np.zeros((shape[0], shape[1], channels))\n",
    "    rstate = np.random.RandomState(seed)\n",
    "    for i in range(0, channels): \n",
    "        \n",
    "   \n",
    "        min_dx = 0\n",
    "        while(min_dx < 3):\n",
    "            xp = np.sort(np.round(rstate.rand(nx-1)*shape[0]).astype(np.int64))\n",
    "            xp = np.concatenate(((0, ), xp, (shape[0], )))\n",
    "            min_dx = np.min(np.diff(xp))\n",
    "        min_dy = 0\n",
    "        while(min_dy < 3):\n",
    "            yp = np.sort(np.round(rstate.rand(ny-1)*shape[1]).astype(np.int64))\n",
    "            yp = np.concatenate(((0, ), yp, (shape[1], )))\n",
    "            min_dy = np.min(np.diff(yp))\n",
    "        img = np.zeros(shape)\n",
    "        for ix, x in enumerate(xp[:-1]):\n",
    "            for iy, y in enumerate(yp[:-1]):\n",
    "                slices = [slice(x, xp[ix+1]), slice(y, yp[iy+1])]\n",
    "                val = rstate.rand(1)[0]\n",
    "                img[slices] = val\n",
    "                \n",
    "        img_final[:,:,i] = img\n",
    "                \n",
    "    return img_final\n",
    "\n",
    "\n",
    "# create an anisotropic piecewise constant image\n",
    "img_mondrian = mondrian((512, 512), channels = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa23e667-8d62-4f41-a27b-8620d018fcc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(img_mondrian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37687127-2db7-401d-89b7-4e44efe47f00",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(r=(0.01/2,0.1,0.01/2))\n",
    "def h(r=.05):\n",
    "    render_display(img_mondrian, r)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addd7125-0c36-41ba-a071-d9d6a7a1a114",
   "metadata": {},
   "source": [
    "----\n",
    "## A Deeper Look at Fourier and Wavelets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae60ff1-f61a-469a-8d81-3b80feed1e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "\n",
    "import numpy as np\n",
    "import scipy as scp\n",
    "import pylab as pyl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from nt_toolbox.general import *\n",
    "from nt_toolbox.signal import *\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c74e4522-bef5-4075-98b4-86a5cf353765",
   "metadata": {},
   "source": [
    "### Signal to Noise Ratio (SNR) \n",
    "*Note:* to measure the error of an image $f$ with its\n",
    "approximation $f_M$, we use the SNR measure, defined as\n",
    "\n",
    "$$ \\text{SNR}(f,f_M) = -20\\log_{10} \\pa{ \\frac{ \\norm{f-f_M} }{ \\norm{f} } },  $$\n",
    "\n",
    "which is a quantity expressed in decibels (dB).\n",
    "The higer the SNR, the better the quality.\n",
    "\n",
    "\n",
    "### Image Loading and Displaying\n",
    "\n",
    "\n",
    "First we load an image $ f \\in \\RR^N $ of $ N = N_0 \\times N_0 $ pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbee53b6-a2b6-4600-ba4a-03587883a7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "n0 = 512\n",
    "f = rescale(load_image(\"nt_toolbox/data/hibiscus.bmp\", n0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9342506c-0468-41e3-82e8-ab111394a78b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(f, 'Image f')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66486087-11b4-479c-a1f9-c32ed123cb35",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(f[n0//2 - 32:n0//2 + 32,n0//2 - 32:n0//2 + 32], 'Zoom')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6801a983-dda3-443a-bba0-46586b948925",
   "metadata": {},
   "source": [
    "An image is a 2D array, it can be modified as a matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e717e0f8-4fda-404c-841f-0f5647517e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "?imageplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a99d4f5-5b0e-4db2-b843-8dc6d01511b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (8,8))\n",
    "imageplot(-f, '-f', [1, 2, 1])\n",
    "imageplot(f[::-1,], 'Flipped', [1, 2, 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e2a95ff-15fb-4749-b431-1c627b911418",
   "metadata": {},
   "source": [
    "Blurring is achieved by computing a convolution $f \\star h$ with a kernel $h$.\n",
    "\n",
    "\n",
    "Compute the low pass kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81aa54ce-4b0c-4944-b745-c44797cd3132",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 9; #size of the kernel\n",
    "h = np.ones([k,k])\n",
    "h = h/np.sum(h) #normalize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b6c1212-a4f3-494a-aff9-97ccf18e76d8",
   "metadata": {},
   "source": [
    "Compute the convolution $f \\star h$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d774110-f009-490d-be70-2eb5790e6033",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import signal\n",
    "fh = signal.convolve2d(f, h, boundary = \"symm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "affd1418-7b9e-4bf8-b1b1-089ab7c140be",
   "metadata": {},
   "outputs": [],
   "source": [
    "fh.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74abff70-8dc2-48ae-84cf-8a30f336f88c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#help(signal.convolve2d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77734bf9-6da0-456f-9c7d-8f147cf7045b",
   "metadata": {},
   "source": [
    "Display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "612a6995-fa93-4c3a-8ea5-a01db66a37e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(fh, 'Blurred image')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d0e9094-e5c6-4a29-8bb1-894a8bd3ee29",
   "metadata": {},
   "source": [
    "Fourier Transform\n",
    "-----------------\n",
    "The Fourier orthonormal basis is defined as\n",
    "$$ \\psi_m(k) = \\frac{1}{\\sqrt{N}}e^{\\frac{2i\\pi}{N_0} \\dotp{m}{k} } $$\n",
    "where $0 \\leq k_1,k_2 < N_0$ are position indexes,\n",
    "and $0 \\leq m_1,m_2 < N_0$ are frequency indexes.\n",
    "\n",
    "\n",
    "The Fourier transform $\\hat f$ is the projection of the image on this Fourier\n",
    "basis\n",
    "\n",
    "$$ \\hat f(m) = \\dotp{f}{\\psi_m}. $$\n",
    "\n",
    "\n",
    "The Fourier transform is computed in $ O(N \\log(N)) $ operation using the FFT algorithm\n",
    "(Fast Fourier Transform).\n",
    "Note the normalization by $\\sqrt{N}=N_0$ to make the transform orthonormal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "effa9900-31b0-432f-b57f-cd76f4ed1fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = pyl.fft2(f)/n0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b305cae6-a537-4033-8a22-f0b0171917d9",
   "metadata": {},
   "source": [
    "Compute the logarithm of the Fourier magnitude $ \\log\\left(\\abs{\\hat f(m)} +\n",
    "\\epsilon\\right) $, for some small $\\epsilon$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b15a8fe-d5fd-45ba-b8bf-ee3f2a559acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "L = pyl.fftshift(np.log(abs(F) + 1e-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74486d8d-7443-47a2-942d-1e382b16e3ba",
   "metadata": {},
   "source": [
    "Display. Note that we use the function fftshift to put the 0 low frequency in the\n",
    "middle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6cc2ae4-c83c-4a1c-ae5a-32a4172042b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(L, 'Log(Fourier transform)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "983a670d-a6cc-4cc7-b420-44bdf3fa390d",
   "metadata": {},
   "source": [
    "Linear Fourier Approximation\n",
    "----------------------------\n",
    "An approximation is obtained by retaining a certain set of index $I_M$\n",
    "\n",
    "$$ f_M = \\sum_{ m \\in I_M } \\dotp{f}{\\psi_m} \\psi_m. $$\n",
    "\n",
    "\n",
    "Linear approximation is obtained by retaining a *fixed* set $I_M$\n",
    "of $M = \\abs{I_M}$ coefficients. The important point is that $I_M$\n",
    "does not depend on the image $f$ to be approximated.\n",
    "\n",
    "\n",
    "For the Fourier transform, a low pass linear approximation\n",
    "is obtained by keeping only the frequencies within a square.\n",
    "\n",
    "$$ I_M = \\enscond{m=(m_1,m_2)}{ -q/2 \\leq m_1,m_2 < q/2 } $$\n",
    "where $ q = \\sqrt{M} $.\n",
    "\n",
    "\n",
    "This can be achieved by computing the Fourier transform,\n",
    "setting to zero the $N-M$ coefficients outside the square $I_M$\n",
    "and then inverting the Fourier transform.\n",
    "\n",
    "\n",
    "Number $M$ of kept coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "673fdc39-bdfc-4a47-b393-6285579f130a",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = n0**2//64"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e70c63bd-7bb3-48fd-935f-58aeffa45a97",
   "metadata": {},
   "source": [
    "\n",
    "### Exercise 1 \n",
    "\n",
    "Perform the linear Fourier approximation with $M$ coefficients.\n",
    "Store the result in the variable $f_M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd8a77a-24ec-4cc6-b796-c16f3211eec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Your Solution Here\n",
    "## Hint: SNR should be 19.8 dB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "655f83c7-93c4-43ac-9d2f-eec843236f42",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "# clamp normalizes the image in [0,1]\n",
    "# snr outputs Signal to noise ration\n",
    "imageplot(clamp(fM), \"Linear, Fourier, SNR = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942e78f3-c98e-4999-a322-cf14c3b62daf",
   "metadata": {},
   "source": [
    "Compare two 1D profile (lines of the image), horizontal first. \n",
    "This shows the strong ringing artifact of the linear approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ecdb788-7cfa-43cb-901c-a74ae874563d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7,6))\n",
    "\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(f[: , n0//2])\n",
    "plt.xlim(0,n0)\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f')\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(fM[: , n0//2])\n",
    "plt.xlim(0,n0)\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f_M')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c383727-6025-476e-9543-d8fd6ec1d65b",
   "metadata": {},
   "source": [
    "Vertical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe8da3f-5c34-4f77-8285-e49a77c77245",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7,6))\n",
    "\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(f[ n0//2, : ])\n",
    "plt.xlim(0,n0)\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f')\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(fM[n0//2,:])\n",
    "plt.xlim(0,n0)\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f_M')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe96e6e-766c-453d-955e-28805bb6431b",
   "metadata": {},
   "source": [
    "Non-linear Fourier Approximation\n",
    "--------------------------------\n",
    "Non-linear approximation is obtained by keeping the $M$ largest\n",
    "coefficients. This is equivalently computed using a thresholding of the\n",
    "coefficients\n",
    "$$ I_M = \\enscond{m}{ \\abs{\\dotp{f}{\\psi_m}}>T }. $$\n",
    "\n",
    "\n",
    "Set a threshold $T>0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5572c50-30c3-43c8-92ad-7769e7fb43f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = .2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc823580-1f3d-48d4-b31c-05b9c19afb4b",
   "metadata": {},
   "source": [
    "Compute the Fourier transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e62317b-00e4-4976-b49f-03a8e11a9af6",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = pyl.fft2(f)/n0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d473fda9-734d-4658-94f2-db205876977a",
   "metadata": {},
   "source": [
    "Do the hard thresholding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb5e54c-1fa4-4f84-a1b3-574ecd195658",
   "metadata": {},
   "outputs": [],
   "source": [
    "FT = np.multiply(F,(abs(F) > T))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70148b46-47e7-43e8-b3f3-37dd1f6e4726",
   "metadata": {},
   "source": [
    "Display. Note that we use the function fftshift to put the 0 low frequency in the\n",
    "middle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "379b1694-9a30-428c-b4c6-ac4b82639688",
   "metadata": {},
   "outputs": [],
   "source": [
    "LT = pyl.fftshift(np.log(abs(FT) + 1e-1))\n",
    "plt.figure(figsize = (10,10))\n",
    "imageplot(L, 'Log(Fourier transform)', sbpt = [1,2,1])\n",
    "imageplot(LT, 'Thresholded Log(Fourier transform)', sbpt = [1, 2, 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6fabff-e5b4-4537-a6a7-60c64d275481",
   "metadata": {},
   "source": [
    "Inverse Fourier transform to obtain $f_M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9e9e224-46c2-4fc5-9b14-ea20652fae78",
   "metadata": {},
   "outputs": [],
   "source": [
    "fM = np.real(pyl.ifft2(FT)*n0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ccf3cd-2355-4874-85eb-8ad8f5f77acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(clamp(fM), \"Linear, Fourier, SNR = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b45bd7f0-b3ca-43d7-a5a0-232ca26aa13e",
   "metadata": {},
   "source": [
    "Given a $T$, the number of coefficients is obtained by counting the\n",
    "non-thresholded coefficients $ \\abs{I_M} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d8aa5d-f72e-4486-9e51-d27fc9937ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = np.sum(FT != 0)\n",
    "print('M/N = 1/%d'  %(n0**2/m))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2573abed-5b4b-41e7-9883-7df57d1b840c",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "Compute the value of the threshold $T$ so that the number of\n",
    "coefficients is $M$. Display the corresponding approximation\n",
    "$f_M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74120152-1f0a-4348-9b04-dbb6f2538031",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Insert your code here.\n",
    "## Hint: SNR should be 20.8 dB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e074c61-20d8-4181-9f5e-105d9f9c02a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "imageplot(clamp(fM), \"Non-linear, Fourier, SNR = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0610e446-a6b9-495c-993d-85b213e1f7c7",
   "metadata": {},
   "source": [
    "Wavelet Transform\n",
    "-----------------\n",
    "A wavelet basis $ \\Bb = \\{ \\psi_m \\}_m $\n",
    "is obtained over the continuous domain by translating and\n",
    "dilating three mother wavelet functions $ \\{\\psi^V,\\psi^H,\\psi^D\\} $.\n",
    "\n",
    "\n",
    "Each wavelet atom is defined as\n",
    "$$ \\psi_m(x) = \\psi_{j,n}^k(x) = \\frac{1}{2^j}\\psi^k\\pa{ \\frac{x-2^j n}{2^j} } $$\n",
    "\n",
    "The scale (size of the support) is $2^j$ and the position is\n",
    "$2^j(n_1,n_2)$. The index is $ m=(k,j,n) $ for $\\{ j \\leq 0 \\}$.\n",
    "\n",
    "\n",
    "The wavelet transform computes all the inner products\n",
    "$ \\{ \\dotp{f}{\\psi_{j,n}^k} \\}_{k,j,n} $.\n",
    "\n",
    "\n",
    "Set the minimum scale for the transform to be 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d5124fe-ba07-46fe-8b99-4f437fe61b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "Jmin = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df315cc9-ccbb-44b8-9bf8-ffab7ea02b3a",
   "metadata": {},
   "source": [
    "Perform the wavelet transform,  $f_w$ stores all the wavelet coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce82f6d-0571-43e7-b35a-41ea8b970271",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nt_toolbox.perform_wavelet_transf import *\n",
    "\n",
    "fw = perform_wavelet_transf(f, Jmin, + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "009a173d-3907-45d9-886d-00594d083a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#?perform_wavelet_transf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a57cf4fd-abae-4eaa-a1e8-ea9610964166",
   "metadata": {},
   "source": [
    "Display the transformed coefficients: all approximations and details coefficients (horizontal, vertical, diagonal)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a7cf90a-c25e-42e1-b873-1ff687683201",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,15))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plot_wavelet(fw)\n",
    "plt.title('Wavelet coefficients')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "imageplot(f, 'Original Image')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4a7ac69-41f2-4587-8ef3-5abbb06f1879",
   "metadata": {},
   "source": [
    "Wavelet Approximation\n",
    "---------------------\n",
    "Linear wavelet approximation with $M=2^{-j_0}$ coefficients is obtained\n",
    "by keeping only the coarse scale (large support) wavelets:\n",
    "\n",
    "$$ I_M = \\enscond{(k,j,n)}{ j \\geq j_0 }. $$\n",
    "\n",
    "\n",
    "It corresponds to setting to zero all the coefficients excepted those\n",
    "that are on the upper left corner of $f_w$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b90ae416-f431-4331-ae5c-bee76c3e5d22",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "Perform linear approximation with $M$ wavelet coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e195b6-b705-4620-87fe-7c1c2930aa5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Insert your code here.\n",
    "## Hint: SNR = 19.8 dB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54f5c8f0-9751-4796-a62d-eeca0282d795",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5,5))\n",
    "imageplot(clamp(fM),  \"Linear, Wavelets, SNR = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2292a416-bf11-4047-9212-ba9886b3131e",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9197f391-6938-4225-bfc7-e9a55040107e",
   "metadata": {},
   "source": [
    "A non-linear approximation is obtained by keeping the $M$ largest wavelet\n",
    "coefficients.\n",
    "\n",
    "\n",
    "As already said, this is equivalently computed by a non-linear hard\n",
    "thresholding.\n",
    "\n",
    "\n",
    "Select a threshold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b0185a8-356a-422e-a8a1-a3c33a74c86c",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = .15"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75d6e578-6ed7-46e7-9d94-9cff27f89269",
   "metadata": {},
   "source": [
    "Perform hard thresholding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc8d581-00f6-4d79-a462-91ed37ba2192",
   "metadata": {},
   "outputs": [],
   "source": [
    "fwT = np.multiply(fw,(abs(fw) > T))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f2d6ec-4f3d-4135-aab3-e29216427746",
   "metadata": {},
   "source": [
    "Display the thresholded coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ed9304-a251-42f5-9b9f-20826108a739",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,15))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plot_wavelet(fw)\n",
    "plt.title('Original coefficients')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plot_wavelet(fwT)\n",
    "plt.title('Thresholded coefficients')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04530381-eda4-452a-a7c8-afccfaf386b9",
   "metadata": {},
   "source": [
    "Perform reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4cbcce0-5855-4b21-af92-04f7783e3840",
   "metadata": {},
   "outputs": [],
   "source": [
    "fM = perform_wavelet_transf(fwT, Jmin, -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0988324-429c-4b13-86f4-ac957f8d3eb1",
   "metadata": {},
   "source": [
    "Display approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43467c3b-2727-4c40-ad15-49428f323186",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5,5))\n",
    "imageplot(clamp(fM), \"Approximation, SNR,  = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c033928-6329-48e0-9287-b748ca3ce581",
   "metadata": {},
   "source": [
    "### Exercise 4 \n",
    "\n",
    "Perform non-linear approximation with $M$ wavelet coefficients by\n",
    "chosing the correct value for $T$. Store the result in the variable\n",
    "$f_M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "225db536-eeb1-44f9-83b8-ad9249f9bef0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Insert your code here.\n",
    "## Hint SNR = 23.8 dB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04f42a0-94cd-4129-8b8e-275c660c9a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5,5))\n",
    "imageplot(clamp(fM), \"Non-linear, Wavelets, SNR = %.1f dB\" %snr(f, fM))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "093e87e0-b033-4f48-8637-241e0ed06f38",
   "metadata": {},
   "source": [
    "Compare two 1D profile (lines of the image).\n",
    "Note how the ringing artifacts are reduced compared to the Fourier\n",
    "approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64fb1b8f-bbc8-4093-82bd-a271351b19e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,8))\n",
    "\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(f[:,n0//2])\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f')\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(fM[:,n0//2])\n",
    "plt.grid(alpha = 0.2)\n",
    "plt.title('f_M')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7573f7bb-4cef-44a4-9a35-726aa4f11a59",
   "metadata": {},
   "source": [
    "## BONUS: Anisotropic Wavelets\n",
    "\n",
    "See \n",
    "https://github.com/PyWavelets/pywt/blob/master/demo/fswavedecn_mondrian.py\n",
    "\n",
    "Using the FSWT to process anistropic images.\n",
    "In this demo, an anisotropic piecewise-constant image is transformed by the\n",
    "standard DWT and the fully-separable DWT. The 'Haar' wavelet gives a sparse\n",
    "representation for such piecewise constant signals (detail coefficients are\n",
    "only non-zero near edges).\n",
    "For such anistropic signals, the number of non-zero coefficients will be lower\n",
    "for the fully separable DWT than for the isotropic one.\n",
    "This example is inspired by the following publication where it is proven that\n",
    "the FSWT gives a sparser representation than the DWT for this class of\n",
    "anistropic images:\n",
    "\n",
    "``` \n",
    "V Velisavljevic, B Beferull-Lozano, M Vetterli and PL Dragotti. Directionlets: Anisotropic Multidirectional Representation With\n",
    "Separable Filtering. IEEE Transactions on Image Processing, Vol. 15, No. 7, July 2006. \n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e030112e-6b53-4143-958b-d1f49e3f26dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#help(pywt.fswavedecn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3caf279-bf98-4161-9155-1a1d4d87ad0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = img_mondrian[:,:,0]\n",
    "# perform DWT\n",
    "coeffs_dwt = pywt.wavedecn(img, wavelet='db1', level=None)\n",
    "\n",
    "# convert coefficient dictionary to a single array\n",
    "coeff_array_dwt, _ = pywt.coeffs_to_array(coeffs_dwt)\n",
    "\n",
    "# perform fully separable DWT\n",
    "fswavedecn_result = pywt.fswavedecn(img, wavelet='db1')\n",
    "\n",
    "nnz_dwt = np.sum(coeff_array_dwt != 0)\n",
    "nnz_fswavedecn = np.sum(fswavedecn_result.coeffs != 0)\n",
    "\n",
    "print(\"Number of nonzero wavedecn coefficients = {}\".format(np.sum(nnz_dwt)))\n",
    "print(\"Number of nonzero fswavedecn coefficients = {}\".format(np.sum(nnz_fswavedecn)))\n",
    "\n",
    "img = mondrian()\n",
    "fig, axes = plt.subplots(1, 3)\n",
    "imshow_kwargs = dict(cmap=plt.cm.gray, interpolation='nearest')\n",
    "axes[0].imshow(img, **imshow_kwargs)\n",
    "axes[0].set_title('Anisotropic Image')\n",
    "axes[1].imshow(coeff_array_dwt != 0, **imshow_kwargs)\n",
    "axes[1].set_title('Nonzero DWT\\ncoefficients\\n(N={})'.format(nnz_dwt))\n",
    "axes[2].imshow(fswavedecn_result.coeffs != 0, **imshow_kwargs)\n",
    "axes[2].set_title('Nonzero FSWT\\ncoefficients\\n(N={})'.format(nnz_fswavedecn))\n",
    "for ax in axes:\n",
    "    ax.set_axis_off()\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}