rkmd.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fmin, fminbound, fmin_bfgs
from scipy.misc import ascent
import pywt
from scipy.stats import kurtosis
true_image = ascent()

sigma = 25.0
test_image = true_image + sigma * np.random.randn(*true_image.shape)

wavelet = 'haar'
levels = 4
coeffs = pywt.wavedec2(test_image, wavelet, level=levels)

lowest_kurtosis = 1e99
for i in range(1, levels + 1)[::-1]:
    hn, vn, dn = coeffs[i]

    def filter(t):
        filtered_coeffs = [pywt.threshold(hn, t),
                           pywt.threshold(vn, t),
                           pywt.threshold(dn, t)]
        return filtered_coeffs

    def restore(t):
        coeffs[i] = filter(t)
        return pywt.waverec2(coeffs, wavelet)

    def objective(t):
        restored = restore(np.abs(t))
        k = kurtosis((test_image - restored).flatten())

        return abs(k) 

    #t_opt = fminbound(objective, 0.0, 5e2, xtol=1e-08)
    t_opt = fmin_bfgs(objective, 10.0, disp=0)
    this_k = objective(t_opt)
    print(i, t_opt, this_k)
    if t_opt > lowest_kurtosis:
        break
    filtered_coeffs = filter(t_opt)
    coeffs[i] = filtered_coeffs

final_restored = pywt.waverec2(coeffs, wavelet)

plt.gray()
plt.subplot(121)
plt.imshow(test_image)
plt.subplot(122)
plt.imshow(final_restored)
plt.show()