"""
Denoising using TOTAL VARIANCE, Works well for random Gaussian noise but not 
as good for salt and pepper. Reference: Sreenivas Bhattiprolu
--------------------------------------------------------------------------------
"""
import cv2, sys, os
from skimage import io, img_as_float
from skimage.restoration import denoise_tv_chambolle
from matplotlib import pyplot as plt

file_name = str(sys.argv[1])
im = os.getcwd() + "\\" + file_name
img = img_as_float(io.imread(im, as_gray=True))

#.flat returns the flattened numpy array (1D)
plt.hist(img.flat, bins=100, range=(0,1))  

"""
denoise_tv_chambolle(
  image, weight=0.1, eps=0.0002, max_num_iter=200, channel_axis=None
)
weight: The greater weight, the more denoising (at expense of fidelity to input)
eps: Delta of the cost function that determines stop criterion
max_num_iter: Maximum number of iterations used for optimization
--------------------------------------------------------------------------------
"""
den_img = denoise_tv_chambolle(
  img, weight=0.1, eps=0.0002, max_num_iter=200, channel_axis=None
)
plt.hist(den_img.flat, bins=100, range=(0,1))  

cv2.imshow("Original", img)
cv2.imshow("Filtered", den_img)
cv2.waitKey(0)          
cv2.destroyAllWindows() 
#-------------------------------------------------------------------------------