diff --git a/functions/NeuralNet_classif.pkl b/functions/NeuralNet_classif.pkl
new file mode 100644
index 0000000..375d597
Binary files /dev/null and b/functions/NeuralNet_classif.pkl differ
diff --git a/functions/sds.py b/functions/sds.py
index e9b73fe..ebcee1b 100644
--- a/functions/sds.py
+++ b/functions/sds.py
@@ -26,7 +26,11 @@ import skimage.transform as transform
 import sklearn.decomposition as decomposition
 import skimage.measure as measure
 import skimage.morphology as morphology
+
+# machine learning modules
 from sklearn.cluster import KMeans
+from sklearn.neural_network import MLPClassifier
+from sklearn.externals import joblib
 
 
 # import own modules
@@ -397,10 +401,6 @@ def pansharpen(im_ms, im_pan, cloud_mask, plot_bool):
     vec_pcs[:,0] = hist_match(vec_pan, vec_pcs[:,0])
     vec_ms_ps = pca.inverse_transform(vec_pcs)
     
-    # normalise between 0 and 1
-    for i in range(vec_pcs.shape[1]):
-        vec_ms_ps[:,i] = np.divide(vec_ms_ps[:,i] - np.min(vec_ms_ps[:,i]),
-                         np.max(vec_ms_ps[:,i]) - np.min(vec_ms_ps[:,i]))
     # reshape vector into image
     vec_ms_ps_full = np.ones((len(vec_mask), im_ms.shape[2])) * np.nan
     vec_ms_ps_full[~vec_mask,:] = vec_ms_ps
@@ -409,11 +409,11 @@ def pansharpen(im_ms, im_pan, cloud_mask, plot_bool):
     if plot_bool:
         plt.figure()
         ax1 = plt.subplot(121)
-        plt.imshow(im_ms[:,:,[2,1,0]])
+        plt.imshow(rescale_image_intensity(im_ms[:,:,[2,1,0]], cloud_mask, 100, False))
         plt.axis('off')
         plt.title('Original')
         ax2 = plt.subplot(122, sharex=ax1, sharey=ax1)
-        plt.imshow(im_ms_ps[:,:,[2,1,0]])
+        plt.imshow(rescale_image_intensity(im_ms_ps[:,:,[2,1,0]], cloud_mask, 100, False))
         plt.axis('off')
         plt.title('Pansharpened')
         plt.show()
@@ -610,6 +610,10 @@ def convert_epsg(points, epsg_in, epsg_out):
 def classify_sand_unsupervised(im_ms_ps, im_pan, cloud_mask, wl_pix, buffer_size, min_beach_size, plot_bool):
     """
     Classifies sand pixels using an unsupervised algorithm (Kmeans)
+    Set buffer size to False if you want to classify the entire image,
+    otherwise buffer size defines the buffer around the shoreline in which 
+    pixels are considered for classification.
+    This classification is not robust and is only used to train a supervised algorithm
     
     KV WRL 2018
 
@@ -623,8 +627,9 @@ def classify_sand_unsupervised(im_ms_ps, im_pan, cloud_mask, wl_pix, buffer_size
             2D cloud mask with True where cloud pixels are
         wl_pix: list of np.ndarray
             list of arrays containig the pixel coordinates of the water line
-        buffer_size: int
+        buffer_size: int or False
             radius of the disk used to create a buffer around the water line
+            when False, the entire image is considered for kmeans
         min_beach_size: int
             minimum number of connected pixels belonging to a single beach
         plot_bool: boolean
@@ -644,16 +649,19 @@ def classify_sand_unsupervised(im_ms_ps, im_pan, cloud_mask, wl_pix, buffer_size
     vec_features[:,[0,1,2,3]] = vec_ms_ps[:,[0,1,2,3]]
     vec_features[:,4] = vec_pan
     
-    # create binary image with ones where the detected water lines is
-    im_buffer = np.zeros((im_ms_ps.shape[0], im_ms_ps.shape[1]))
-    for i, contour in enumerate(wl_pix):
-            indices = [(int(_[0]), int(_[1])) for _ in list(np.round(contour))]
-            for j, idx in enumerate(indices):
-                im_buffer[idx] = 1
-    # perform a dilation on the binary image
-    se = morphology.disk(buffer_size)
-    im_buffer = morphology.binary_dilation(im_buffer, se)
-    vec_buffer = (im_buffer == 1).reshape(im_ms_ps.shape[0] * im_ms_ps.shape[1])
+    if buffer_size:
+        # create binary image with ones where the detected water lines is
+        im_buffer = np.zeros((im_ms_ps.shape[0], im_ms_ps.shape[1]))
+        for i, contour in enumerate(wl_pix):
+                indices = [(int(_[0]), int(_[1])) for _ in list(np.round(contour))]
+                for j, idx in enumerate(indices):
+                    im_buffer[idx] = 1
+        # perform a dilation on the binary image
+        se = morphology.disk(buffer_size)
+        im_buffer = morphology.binary_dilation(im_buffer, se)
+        vec_buffer = (im_buffer == 1).reshape(im_ms_ps.shape[0] * im_ms_ps.shape[1])
+    else:
+        vec_buffer = np.ones((vec_pan.shape[0]))
     # add cloud mask to buffer
     vec_buffer= np.logical_and(vec_buffer, ~vec_mask)
     # perform kmeans (6 clusters)
@@ -664,17 +672,98 @@ def classify_sand_unsupervised(im_ms_ps, im_pan, cloud_mask, wl_pix, buffer_size
     # find the class with maximum reflection in the B,G,R,Pan
     im_sand = im_labels == np.argmax(np.mean(kmeans.cluster_centers_[:,[0,1,2,4]], axis=1))
     im_sand = morphology.remove_small_objects(im_sand, min_size=min_beach_size, connectivity=2)
+    im_sand = morphology.binary_erosion(im_sand, morphology.disk(1))
 #    im_sand = morphology.binary_dilation(im_sand, morphology.disk(1))
     
     if plot_bool:
-        im = np.copy(im_ms_ps)
+        im = np.copy(rescale_image_intensity(im_ms_ps[:,:,[2,1,0]], cloud_mask, 100, False))
         im[im_sand,0] = 0
         im[im_sand,1] = 0
         im[im_sand,2] = 1
         plt.figure()
-        plt.imshow(im[:,:,[2,1,0]])
+        plt.imshow(im)
         plt.axis('image')
         plt.title('Sand classification')
         plt.show()
     
     return im_sand
+
+def classify_image_NN(im_ms_ps, im_pan, cloud_mask, plot_bool):
+    """
+    Classifies every pixel in the image in one of 4 classes:
+        - sand                                          --> label = 1
+        - whitewater (breaking waves and swash)         --> label = 2
+        - water                                         --> label = 3
+        - other (vegetation, buildings, rocks...)       --> label = 0
+    
+    The classifier is a Neural Network, trained with 7000 pixels for the class SAND and 1500 pixels for
+    each of the other classes. This is because the class of interest for my application is SAND and I 
+    wanted to minimize the classification error for that class
+    
+    KV WRL 2018
+
+    Arguments:
+    -----------
+        im_ms_ps: np.ndarray
+            Pansharpened RGB + downsampled NIR and SWIR
+        im_pan:
+            Panchromatic band
+        cloud_mask: np.ndarray
+            2D cloud mask with True where cloud pixels are
+        plot_bool: boolean
+            True if plot is wanted
+                
+    Returns:    -----------
+        im_labels: np.ndarray
+            2D binary image containing True where sand pixels are located
+
+    """     
+    
+    # load classifier
+    clf = joblib.load('functions/NeuralNet_classif.pkl')
+    
+    # calculate features
+    n_features = 10
+    im_features = np.zeros((im_ms_ps.shape[0], im_ms_ps.shape[1], n_features))
+    im_features[:,:,[0,1,2,3,4]] = im_ms_ps
+    im_features[:,:,5] = im_pan
+    im_features[:,:,6] = nd_index(im_ms_ps[:,:,3], im_ms_ps[:,:,1], cloud_mask, False) # (NIR-G)
+    im_features[:,:,7] = nd_index(im_ms_ps[:,:,3], im_ms_ps[:,:,2], cloud_mask, False) # ND(NIR-R)
+    im_features[:,:,8] = nd_index(im_ms_ps[:,:,0], im_ms_ps[:,:,2], cloud_mask, False) # ND(B-R)
+    im_features[:,:,9] = nd_index(im_ms_ps[:,:,4], im_ms_ps[:,:,1], cloud_mask, False) # ND(SWIR-G)
+    # remove NaNs and clouds
+    vec_features = im_features.reshape((im_ms_ps.shape[0] * im_ms_ps.shape[1], n_features))
+    vec_cloud = cloud_mask.reshape(cloud_mask.shape[0]*cloud_mask.shape[1])
+    vec_nan = np.any(np.isnan(vec_features), axis=1)
+    vec_mask = np.logical_or(vec_cloud, vec_nan)
+    vec_features = vec_features[~vec_mask, :]
+    # predict with NN classifier
+    labels = clf.predict(vec_features)
+    # recompose image
+    vec_classif = np.zeros((cloud_mask.shape[0]*cloud_mask.shape[1])) 
+    vec_classif[~vec_mask] = labels
+    im_classif = vec_classif.reshape((im_ms_ps.shape[0], im_ms_ps.shape[1]))
+#    im_classif = morphology.remove_small_objects(im_classif, min_size=min_beach_size, connectivity=2)
+        
+    if plot_bool:
+        im_display = rescale_image_intensity(im_ms_ps[:,:,[2,1,0]], cloud_mask, 100, False)
+        im = np.copy(im_display)
+        colours = np.array([[1,0,0],[1,1,0],[0,1,1],[0,0,1]])
+        for k in range(4):
+            im[im_classif == k,0] = colours[k,0]
+            im[im_classif == k,1] = colours[k,1]
+            im[im_classif == k,2] = colours[k,2]
+        
+        plt.figure()
+        ax1 = plt.subplot(121)
+        plt.imshow(im_display)
+        plt.axis('off')
+        plt.title('Image')
+        ax2 = plt.subplot(122, sharex=ax1, sharey=ax1)
+        plt.imshow(im)
+        plt.axis('off')
+        plt.title('NN')  
+        plt.draw()
+    
+    return im_classif
+
diff --git a/sds_extract.py b/sds_extract.py
index a6f4126..74742c3 100644
--- a/sds_extract.py
+++ b/sds_extract.py
@@ -37,7 +37,6 @@ ee.Initialize()
 # parameters
 cloud_thresh = 0.5      # threshold for cloud cover
 plot_bool = True      # if you want the plots
-prob_high = 99.9        # upper probability to clip and rescale pixel intensity
 min_contour_points = 100# minimum number of points contained in each water line
 output_epsg = 28356     # GDA94 / MGA Zone 56
 buffer_size = 10        # radius of disk for buffer (sand classif parameter)
@@ -71,23 +70,26 @@ i = 0 # first image
 im = ee.Image(im_all[i].get('id')) 
 
 # load image as np.array
-im_pan, im_ms, im_cloud, crs, meta = sds.read_eeimage(im, polygon, satname, plot_bool)
+im_pan, im_ms, cloud_mask, crs, meta = sds.read_eeimage(im, polygon, satname, plot_bool)
 
-# rescale intensities
-im_ms = sds.rescale_image_intensity(im_ms, im_cloud, prob_high, plot_bool)
-im_pan = sds.rescale_image_intensity(im_pan, im_cloud, prob_high, plot_bool)
+# mask -inf or nan values on the image and add to cloud_mask
+im_inf = np.isin(im_ms[:,:,0], -np.inf)
+im_nan = np.isnan(im_ms[:,:,0])
+cloud_mask = np.logical_or(np.logical_or(cloud_mask, im_inf), im_nan)
+cloud_cover = sum(sum(cloud_mask.astype(int)))/(cloud_mask.shape[0]*cloud_mask.shape[1])
+print('Cloud cover : ' + str(int(round(100*cloud_cover))) + ' %')
 
 # pansharpen rgb image
-im_ms_ps = sds.pansharpen(im_ms[:,:,[0,1,2]], im_pan, im_cloud, plot_bool)
+im_ms_ps = sds.pansharpen(im_ms[:,:,[0,1,2]], im_pan, cloud_mask, plot_bool)
 
 # add down-sized bands for NIR and SWIR (since pansharpening is not possible)
 im_ms_ps = np.append(im_ms_ps, im_ms[:,:,[3,4]], axis=2) 
 
 # calculate NDWI
-im_ndwi = sds.nd_index(im_ms_ps[:,:,3], im_ms_ps[:,:,1], im_cloud, plot_bool)
+im_ndwi = sds.nd_index(im_ms_ps[:,:,3], im_ms_ps[:,:,1], cloud_mask, plot_bool)
 
 # edge detection
-wl_pix = sds.find_wl_contours(im_ndwi, im_cloud, min_contour_points, plot_bool)
+wl_pix = sds.find_wl_contours(im_ndwi, cloud_mask, min_contour_points, plot_bool)
 
 plt.figure()
 plt.imshow(im_ms_ps[:,:,[2,1,0]])
@@ -102,5 +104,8 @@ wl_coords = sds.convert_pix2world(wl_pix, crs['crs_15m'])
 # convert to output epsg spatial reference
 wl = sds.convert_epsg(wl_coords, crs['epsg_code'], output_epsg)
 
-# classify sand pixels
-im_sand = sds.classify_sand_unsupervised(im_ms_ps, im_pan, im_cloud, wl_pix, buffer_size, min_beach_size, plot_bool)
+# classify sand pixels with Kmeans
+#im_sand = sds.classify_sand_unsupervised(im_ms_ps, im_pan, cloud_mask, wl_pix, buffer_size, min_beach_size, plot_bool)
+
+# classify image in 4 classes (sand, whitewater, water, other) with NN classifier
+im_classif = sds.classify_image_NN(im_ms_ps, im_pan, cloud_mask, plot_bool)
\ No newline at end of file