denose with dwt

import sys
import cv2
import numpy as np
import pywt

def apply_thresh(coef, thresh, mode):
	if mode == 'hard':
		tmp = np.abs(coef)
		tmp = tmp >= thresh
		coef = coef * tmp
	else: #soft
		tmp = np.abs(coef)
		tmp = tmp >= thresh
		coef = coef * tmp
		coef = np.sign(coef) * (np.abs(coef) - thresh)
	return coef

def get_thresh_uplimit(coef):
	c = 13 #noise estimation
	s = coef.shape[0] * coef.shape[1]
	return c * np.sqrt(2 * np.log(s))

def get_thresh_MAD(coef): #mean absolute difference
	c = 2.0
	coef = np.abs(coef)
	return c * np.std(np.reshape(coef, coef.shape[0]*coef.shape[1]))
	
def denoise_dwt(img):
	if img.ndim != 2:
		img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
	img = np.float32(img)
	h,w = img.shape
	if h < 32 or w < 32:
		return np.zeros(img.shape,np.uint8)
	get_thresh = get_thresh_uplimit
	N = 3
	coef = pywt.wavedec2(img, 'db1', level = N)
	for k in range(2,N+1):
		cH, cV, cD = coef[k]
		t = get_thresh(cH)
		cH = apply_thresh(cH, t, 'soft')
		t = get_thresh(cV)
		cV = apply_thresh(cV, t, 'soft')
		t = get_thresh(cD)
		cD = apply_thresh(cD, t, 'soft')
		coef[k] = (cH, cV, cD)
	cleanimg = pywt.waverec2(coef, 'db1')
	cleanimg = np.uint8(cleanimg)
	return cleanimg

def imnoise(img):
	if img.ndim != 2:
		img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
	noiseimg = np.float32(img) + np.random.normal(0, 50, img.shape)
	m0 = np.zeros(img.shape, float)
	m1 = np.ones(img.shape,float) * 255
	noiseimg = np.minimum(noiseimg, m1)
	noiseimg = np.maximum(noiseimg, m0)
	noiseimg = np.uint8(noiseimg)
	return noiseimg

if __name__=='__main__':
	img = cv2.imread(sys.argv[1])
	cv2.imshow('src', img)
	cv2.waitKey(-1)
	noiseimg = imnoise(img)
	cv2.imshow('noise', noiseimg)
	cv2.waitKey(-1)
	cleanimg = denoise_dwt(noiseimg)
	cv2.imshow('clean', cleanimg)
	cv2.waitKey(-1)