from skimage import io, img_as_float, exposure, data
from skimage.util.dtype import dtype_range
from scipy.ndimage import median_filter, gaussian_filter, binary_fill_holes

import glob
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline

def plot_img_and_hist(img, axes, bins=256):
    """Plot an image along with its histogram and cumulative histogram.

    """
    ax_img, ax_hist = axes
    ax_cdf = ax_hist.twinx()

    # Display image
    ax_img.imshow(img, cmap=plt.cm.gray)
    ax_img.set_axis_off()

    # Display histogram
    ax_hist.hist(img.ravel(), bins=bins)
    ax_hist.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
    ax_hist.set_xlabel('Pixel intensity')

    xmin, xmax = dtype_range[img.dtype.type]
    ax_hist.set_xlim(xmin, xmax)

    # Display cumulative distribution
    img_cdf, bins = exposure.cumulative_distribution(img, bins)
    ax_cdf.plot(bins, img_cdf, 'r')

    return ax_img, ax_hist, ax_cdf

def dynamic_masking(image,threshold = 0.4, filter_size=7,filter_sigma=5):
    """ Dynamically masks out the objects in the PIV images
    """
    image = exposure.rescale_intensity(image, in_range=(0, 1))
    blurback = gaussian_filter(median_filter(image,size=filter_size),sigma=filter_sigma)
    # create the boolean mask 
    bw = (blurback > threshold).astype('bool')
    bw = binary_fill_holes(bw)
    image[bw] = 0.0
    # image -= blurback
    image = exposure.rescale_intensity(image,in_range=(0,1))
    return image.copy()

# Example:
imdir = '/Volumes/ALEX/Mark_Exp_01-06-14/water/10_53Hz_3/'
filelist = sorted(glob.glob(imdir+'*.tif'))
image = img_as_float(io.imread(filelist[1]))
masked_image = dynamic_masking(image)
f, axes = plt.subplots(2, 2, figsize=(15, 15))
ax_img, ax_hist, ax_cdf = plot_img_and_hist(image, axes[:, 0])
ax_img.set_title('Original PIV image')
ax_img, ax_hist, ax_cdf = plot_img_and_hist(masked_image, axes[:, 1])
ax_img.set_title('Masked PIV image')
plt.subplots_adjust(wspace=0.4)
plt.show()


plt.figure(figsize=(10,10))
plt.imshow(np.c_[image,dynamic_masking(image,threshold = 0.45, filter_size=5,filter_sigma=3)],cmap=plt.cm.gray)

plt.figure(figsize=(10,10))
plt.imshow(np.c_[image,dynamic_masking(image,threshold = 0.45, filter_size=15,filter_sigma=5.0)],cmap=plt.cm.gray)