#!/usr/bin/env python
# coding: utf-8

# Wavelet Block Thresholding
# ==========================
# 
# *Important:* Please read the [installation page](http://gpeyre.github.io/numerical-tours/installation_python/) for details about how to install the toolboxes.
# $\newcommand{\dotp}[2]{\langle #1, #2 \rangle}$
# $\newcommand{\enscond}[2]{\lbrace #1, #2 \rbrace}$
# $\newcommand{\pd}[2]{ \frac{ \partial #1}{\partial #2} }$
# $\newcommand{\umin}[1]{\underset{#1}{\min}\;}$
# $\newcommand{\umax}[1]{\underset{#1}{\max}\;}$
# $\newcommand{\umin}[1]{\underset{#1}{\min}\;}$
# $\newcommand{\uargmin}[1]{\underset{#1}{argmin}\;}$
# $\newcommand{\norm}[1]{\|#1\|}$
# $\newcommand{\abs}[1]{\left|#1\right|}$
# $\newcommand{\choice}[1]{ \left\{  \begin{array}{l} #1 \end{array} \right. }$
# $\newcommand{\pa}[1]{\left(#1\right)}$
# $\newcommand{\diag}[1]{{diag}\left( #1 \right)}$
# $\newcommand{\qandq}{\quad\text{and}\quad}$
# $\newcommand{\qwhereq}{\quad\text{where}\quad}$
# $\newcommand{\qifq}{ \quad \text{if} \quad }$
# $\newcommand{\qarrq}{ \quad \Longrightarrow \quad }$
# $\newcommand{\ZZ}{\mathbb{Z}}$
# $\newcommand{\CC}{\mathbb{C}}$
# $\newcommand{\RR}{\mathbb{R}}$
# $\newcommand{\EE}{\mathbb{E}}$
# $\newcommand{\Zz}{\mathcal{Z}}$
# $\newcommand{\Ww}{\mathcal{W}}$
# $\newcommand{\Vv}{\mathcal{V}}$
# $\newcommand{\Nn}{\mathcal{N}}$
# $\newcommand{\NN}{\mathcal{N}}$
# $\newcommand{\Hh}{\mathcal{H}}$
# $\newcommand{\Bb}{\mathcal{B}}$
# $\newcommand{\Ee}{\mathcal{E}}$
# $\newcommand{\Cc}{\mathcal{C}}$
# $\newcommand{\Gg}{\mathcal{G}}$
# $\newcommand{\Ss}{\mathcal{S}}$
# $\newcommand{\Pp}{\mathcal{P}}$
# $\newcommand{\Ff}{\mathcal{F}}$
# $\newcommand{\Xx}{\mathcal{X}}$
# $\newcommand{\Mm}{\mathcal{M}}$
# $\newcommand{\Ii}{\mathcal{I}}$
# $\newcommand{\Dd}{\mathcal{D}}$
# $\newcommand{\Ll}{\mathcal{L}}$
# $\newcommand{\Tt}{\mathcal{T}}$
# $\newcommand{\si}{\sigma}$
# $\newcommand{\al}{\alpha}$
# $\newcommand{\la}{\lambda}$
# $\newcommand{\ga}{\gamma}$
# $\newcommand{\Ga}{\Gamma}$
# $\newcommand{\La}{\Lambda}$
# $\newcommand{\si}{\sigma}$
# $\newcommand{\Si}{\Sigma}$
# $\newcommand{\be}{\beta}$
# $\newcommand{\de}{\delta}$
# $\newcommand{\De}{\Delta}$
# $\newcommand{\phi}{\varphi}$
# $\newcommand{\th}{\theta}$
# $\newcommand{\om}{\omega}$
# $\newcommand{\Om}{\Omega}$

# This numerical tour presents block thresholding methods,
# that makes use of the structure of wavelet coefficients of natural images to perform denoising.
# Theoretical properties of block thresholding were investigated
# in [CaiSilv](#biblio) [Cai99](#biblio) [HallKerkPic99](#biblio)

# In[44]:


from __future__ import division

import numpy as np
import scipy as scp
import pylab as pyl
import matplotlib.pyplot as plt

from nt_toolbox.general import *
from nt_toolbox.signal import *

import warnings
warnings.filterwarnings('ignore')

get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')


# Generating a Noisy Image
# ------------------------
# Here we use an additive Gaussian noise.
# 
# 
# Size of the image of $N=n \times n$ pixels.

# In[45]:


n = 256


# First we load an image $f_0 \in \RR^N$.

# In[46]:


f0 = rescale(load_image("nt_toolbox/data/boat.bmp", n))


# Display it.

# In[47]:


plt.figure(figsize = (5,5))
imageplot(f0)


# Noise level.

# In[48]:


sigma = .08


# Generate a noisy image $f=f_0+\epsilon$ where $\epsilon \sim
# \Nn(0,\si^2\text{Id}_N)$.

# In[49]:


from numpy import random

f = f0 + sigma*random.standard_normal((n,n))


# Display it.

# In[50]:


plt.figure(figsize = (5,5))
imageplot(clamp(f))


# Orthogonal Wavelet Thresholding
# -------------------------------
# We first consider the traditional
# wavelet thresholding method.
# 
# 
# Parameters for the orthogonal wavelet transform.

# In[51]:


Jmin = 4


# Shortcuts for the foward and backward wavelet transforms.

# In[52]:


from nt_toolbox.perform_wavelet_transf import *

wav  = lambda f  : perform_wavelet_transf(f, Jmin, +1)
iwav = lambda fw : perform_wavelet_transf(fw, Jmin, -1)


# Display the original set of noisy coefficients.

# In[53]:


plt.figure(figsize=(10,10))
plot_wavelet(wav(f), Jmin)
plt.show()


# Denoting $\Ww$ and $\Ww^*$ the forward and backward wavelet
# transform, wavelet thresholding $\tilde f$ is defined as
# 
# $$ \tilde f = \Ww^* \circ \theta_T \circ \Ww(f) $$
# 
# where $T>0$ is the threshold, that should be adapted to the noise
# level.
# 
# 
# The thresholding operator is applied component-wise
# 
# $$ \th_T(x)_i = \psi_T(x_i) x_i $$
# 
# where $\psi_T$ is an atenuation fonction. In this tour, we use the James
# Stein (JS) attenuation:
# 
# $$ \psi_T(s) = \max\pa{ 0, 1-\frac{T^2}{s^2} } $$

# In[54]:


psi = lambda s,T : np.maximum(1-T**2/np.maximum(abs(s)**2, 1e-9*np.ones(np.shape(s))), np.zeros(np.shape(s)))


# Display the thresholding function $\th_T$.

# In[55]:


s = np.linspace(-3, 3, 1024)

plt.plot(s, s*psi(s,1))
plt.plot(s, s, 'r--')

plt.show()


# Thresholding operator.

# In[56]:


theta = lambda x,T : psi(x, T)*x
ThreshWav = lambda f,T : iwav(theta(wav(f), T))


# Test the thresholding.

# In[57]:


T = 1.5*sigma

plt.figure(figsize=(5,5))
imageplot(clamp(ThreshWav(f, T)))


# __Exercise 1__
# 
# Display the evolution of the denoising SNR when $T$ varies.
# Store in $f_{Thresh}$ the optimal denoising result.

# In[58]:


run -i nt_solutions/denoisingwav_4_block/exo1


# In[59]:


## Insert your code here.


# Display the optimal thresolding.

# In[60]:


plt.figure(figsize = (5,5))
imageplot(clamp(fThresh), "SNR = %.1f dB" %snr(f0, fThresh))


# Block Thresholding Operator
# ---------------------------
# A block thresholding operator of coefficients $x=(x_i)_{i=1}^P \in \RR^P$ is defined
# using a partition $B$ into a set of blocks $b$
# 
# $$ \{1,\ldots,P\} = \bigcup_{b \in B} b. $$
# Its definition reads
# 
# $$ \forall i \in b, \quad
#       \theta_T(x)_i = \psi_T\left( \norm{x_b}_2 \right) x_i $$
#       
# where $ x_b = (x_j)_{j \in B} \in \RR^{\abs{b}} $.
# One thus thresholds the $\ell^2$ norm (the energy) of each block rather
# than each coefficient independently.
# 
# 
# For image-based thresholding, we use a partition in square blocks of
# equal size $w \times w$.
# 
# 
# The block size $w$.

# In[61]:


w = 4
n = 256


# Compute indexing of the blocks.

# In[62]:


[X,Y,dX,dY] = np.meshgrid(np.arange(1,n-w+2,w),np.arange(1,n-w+2,w),np.arange(0,w),np.arange(0,w))
I = (X + dX-1) + (Y + dY-1)*n
for k in range(n//w):
    for l in range(n//w):
        I[k][l] = np.transpose(I[k][l])


# Block extraction operator. It returns the set $ \{x_b\}_{b \in B} $
# of block-partitioned coefficients.

# In[63]:


block = lambda x : np.ravel(x)[I]


# Block reconstruction operator.

# In[64]:


def assign(M,I,H):
    M_temp = M
    np.ravel(M_temp)[I] = H
    return np.reshape(M_temp,(n,n))

iblock = lambda H : assign(np.zeros([n,n]), I, H)


# Check that block extraction / reconstruction gives perfect
# reconstruction.

# In[65]:


from numpy import linalg

print("Should be 0:", linalg.norm(f - iblock(block(f))))


# Compute the average energy of each block, and duplicate.

# In[66]:


def energy(H):
    H_tmp = np.copy(H)
    for i in range(n//w):
        for j in range(n//w):
            H_tmp[i][j] = np.sqrt(np.mean(H_tmp[i][j]**2))#*np.ones([1,1])
    return H_tmp


# Block thresholding operator.

# In[67]:


Thresh = lambda H,T : psi(energy(H), T)*H
ThreshBlock = lambda x,T : iblock(Thresh(block(x), T))


# __Exercise 2__
# 
# Test the effect of block thresholding on the image $f_0$ itself, for increasing value of $T$.
# Of course directly thresholding the image has no interest, this is just
# to vizualize the effect.

# In[68]:


run -i nt_solutions/denoisingwav_4_block/exo2


# In[69]:


## Insert your code here.


# Orthogonal Wavelet Block Thresholding
# -------------------------------------
# Wavelet coefficients of natural images are not independant one from each
# other. One can thus improve the denoising results by thresholding block
# of coefficients togethers. Block thresholding is only efficient when
# used as a soft thresholder. Here we use a Stein soft thresholder.
# 
# 
# Display the thresholded coefficients for a threshold value $T$ proportional to the noise level $\si$.

# In[70]:


T = 1.25*sigma

plt.figure(figsize=(10,10))
plot_wavelet(ThreshBlock(wav(f), T), Jmin)
plt.show()


# Define the wavelet block thresholding operator.

# In[71]:


ThreshWav = lambda f, T: iwav(ThreshBlock(wav(f), T))


# Test the thresholding.

# In[72]:


plt.figure(figsize=(5,5))
imageplot(clamp(ThreshWav(f, T)))


# __Exercise 3__
# 
# Display the evolution of the denoising SNR when $T$ varies.
# Store the optimal denoising result in $f_{Block}$.

# In[73]:


run -i nt_solutions/denoisingwav_4_block/exo3


# In[74]:


## Insert your code here.


# Display the result.

# In[75]:


plt.figure(figsize=(5,5))
imageplot(clamp(fBlock), "SNR = %.1f dB" %snr(f0, fBlock))


# Translation invariant Block Thresholding
# ----------------------------------------
# Block thresholding can also be applied to a translation invariant wavelet
# transform. It gives state of the art denoising results.
# 
# 
# 
# Shortcuts for the foward and backward translation invariant wavelet transforms.

# In[76]:


wav  = lambda f: perform_wavelet_transf(f, Jmin, + 1, ti=1)
iwav = lambda fw: perform_wavelet_transf(fw, Jmin, -1, ti=1)


# Foward wavelet transform.

# In[77]:


fw = wav(f)
n = 256
#np.shape(fw)[0]


# Compute indexing of the blocks.

# In[78]:


[X,J,Y,dX,dY] = np.meshgrid(np.arange(1,n-w+2,w),np.arange(1,np.shape(fw)[0]+1),np.arange(1,n-w+2,w),np.arange(0,w),np.arange(0,w))
I = (X + dX-1) + (Y + dY-1)*n + (J-1)*n**2
for k in range(n//w):
    for l in range(n//w):
        for m in range(np.shape(fw)[0]):
            I[m][k][l] = np.transpose(I[m][k][l])


# Forward and backward extraction operators.

# In[79]:


block = lambda x : np.ravel(x)[I]

def assign(M,I,H):
    M_temp = M
    np.ravel(M_temp)[I] = H
    return np.reshape(M_temp,(np.shape(fw)[0],n,n))

iblock = lambda H : assign(np.zeros([np.shape(fw)[0],n,n]), I, H)


# Compute the average energy of each block, and duplicate.

# In[80]:


def energy(H):
    H_tmp = np.copy(H)
    for i in range(n//w):
        for j in range(n//w):
            for k in range(np.shape(fw)[0]):
                    H_tmp[k][i][j] = np.sqrt(np.mean(H_tmp[k][i][j]**2))
    return H_tmp


# Block thresholding operator.

# In[81]:


Thresh = lambda H,T : psi(energy(H), T)*H
ThreshBlock = lambda x,T : iblock(Thresh(block(x), T))


# Define the wavelet block thresholding operator.

# In[82]:


ThreshWav = lambda f,T : iwav(ThreshBlock(wav(f), T))


# Test the thresholding.

# In[ ]:


T = 1.25*sigma

plt.figure(figsize = (5,5))
imageplot(clamp(ThreshWav(f, T)))


# __Exercise 4__
# 
# Display the evolution of the denoising SNR when $T$ varies.
# Store the optimal denoising result in $f_{TI}$.

# In[ ]:


run -i nt_solutions/denoisingwav_4_block/exo4


# In[ ]:


## Insert your code here.


# Display the result.

# In[ ]:


plt.figure(figsize = (5,5))
imageplot(clamp(fTI), "SNR = %.1f dB" %snr(f0, fTI))


# Bibliography
# ------------
# <html><a name="biblio"></a></html>
# 
# 
# * [CaiSil01] T. Cai and B.W. Silverman, [Incorporating information on neighboring coefficients into wavelet estimation][1], Sankhya 63, 127-148, 2001.
# * [Cai99] T. Cai, [Adaptive wavelet estimation: a block thresholding and oracle inequality approach][2], The Annals of Statistics 27, 898-924, 1999.
# * [HallKerkPic99] P. Hall, G. Kerkyacharian and D. Picard, _On the minimax optimality of block thresholded wavelet estimator_, Statistica Sinica 9(1999), 33-49
# 
# [1]:http://sankhya.isical.ac.in/search/63b2/caifnl.html
# [2]:http://dx.doi.org/10.1214/aos/1018031262
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>