Added functionality to mask by land cover maps

biota/core.py

@@ -349,17 +349,30 @@ def getMask(self, masked_px_count = False):
 
         return mask
 
-    def updateMask(self, shp, buffer_size = 0.):
+    def updateMask(self, filename, buffer_size = 0., classes = []):
         """
-        Function to add further pixels to mask based on shapefiles and buffers.
+        Function to add further pixels to mask based on a shapefile or GeoTiff file, with optional buffers.
         """
 
-        # Rasterize the shapefile, optionally with a buffer
-        shp_mask = biota.mask.rasterizeShapefile(self, shp, buffer_size = buffer_size)
+        file_type = filename.split('/')[-1].split('.')[-1]
 
-        # Add the rasterized shapefile to the mask
-        self.mask = np.logical_or(self.mask, shp_mask)     
+        assert file_type in ['shp', 'tif', 'tiff'], "Input filename must be a GeoTiff or a shapefile."
+
+        if file_type == 'shp':
+
+            # Rasterize the shapefile, optionally with a buffer
+            mask = biota.mask.rasterizeShapefile(self, filename, buffer_size = buffer_size)
+
+        if file_type in ['tif', 'tiff']:
+
+            assert classes != [], "If adding a GeoTiff file to the mask, you must also specify the class values to add to the mask (e.g. classes = [20, 160, 170, 190, 210])."
+
+            # Resample and extract values from shapefile, optionally with a buffer
+            mask = biota.mask.maskRaster(self, filename, classes = classes, buffer_size = buffer_size)
 
+        # Add the new raster masks to the existing mask
+        self.mask = np.logical_or(self.mask, mask)
+
     def resetMask(self):
         """
         Function to reset a mask to the default.

biota/mask.py

@@ -6,6 +6,8 @@
 from PIL import Image, ImageDraw
 import scipy.ndimage as ndimage
 
+import biota.IO
+
 import pdb
 
 
@@ -162,6 +164,52 @@ def getBBox(shp, field, value):
     return bbox
 
 
+def maskRaster(data, tif, classes = [], buffer_size = 0.):
+    '''
+    Extract a mask from a GeoTiff based on specified classes
+    
+    Args:
+        data: An ALOS object
+        tif: A GeoTiff file, with integer values
+        classes: A list of values to add to be masked
+        buffer_size: Optionally specify a buffer to add around maked pixels, in meters.
+    
+    Returns:
+        A numpy array with a boolean mask
+    '''
+
+    from osgeo import gdal
+
+    assert tif.rstrip('/').split('.')[-1] == 'tif' or tif.rstrip('/').split('.')[-1] == 'tiff', "tif input must be a GeoTiff file."
+
+    # Open GeoTiff and get metadata
+    ds_source = gdal.Open(tif)
+    proj_source = biota.IO.loadProjection(tif)
+
+    # Create output file matching ALOS tile
+    gdal_driver = gdal.GetDriverByName('MEM')
+    ds_dest = gdal_driver.Create('', data.ySize, data.xSize, 1, 3)
+    ds_dest.SetGeoTransform(data.geo_t)
+    ds_dest.SetProjection(data.proj)
+
+    # Reproject input GeoTiff to match the ALOS tile
+    gdal.ReprojectImage(ds_source, ds_dest, proj_source, data.proj, gdal.GRA_NearestNeighbour)
+
+    # Load resampled image into memory
+    tif_resampled = ds_dest.GetRasterBand(1).ReadAsArray()
+
+    # Identify pixels that are classes to be masked
+    mask = np.in1d(tif_resampled, classes).reshape(tif_resampled.shape)
+
+    if buffer_size > 0.:
+
+        # Determine the size of the buffer in pixels
+        buffer_px = int(buffer_size / ((data.xRes + data.yRes) / 2.))
+
+        # Dilate the mask
+        mask = dilateMask(mask, buffer_px)
+
+    return mask
 
 
 def rasterizeShapefile(data, shp, buffer_size = 0., field = None, value = None, location_id = False):
@@ -282,8 +330,6 @@ def rasterizeShapefile(data, shp, buffer_size = 0., field = None, value = None,
     return mask
 
 
-
-
 def getTilesInShapefile(shp, field = None, value = None):
     """
     Identify all the ALOS tiles that fall within a shapefile.