Add option to return a rasterized cloud mask

true-color-image/sentinel_helpers.py

@@ -2,10 +2,12 @@ import urllib.parse
 from pathlib import Path
 
 import fiona
+import folium
 import geopandas as gpd
 from matplotlib import pyplot as plt
 import numpy as np
 import rasterio as r
+from rasterio.features import geometry_mask
 from rasterio.warp import calculate_default_transform, reproject, Resampling
 
 from shapely.geometry import shape
@@ -109,11 +111,17 @@ def scihub_bgr_paths(product_path, resolution=None):
     return scihub_band_paths(product_path, ['B02', 'B03', 'B04'], resolution)
 
 
-def scihub_cloud_mask(product_path):
+def scihub_cloud_mask(product_path, **kwargs):
     '''
     Given a `product_path` pointing to a product downlaoded from the Copernicus
     Open Access Hub, returns a shapely geometry representing the included cloud
     mask.
+    
+    If an additional parameter, `rasterize=True` is given, the returned cloud
+    mask will be a rasterized numpy ndarray instead of a vector geometry. Two
+    additional parameters, `target_path` and `target_transform` are needed to
+    determine the size of this array. In this array, pixels with clouds are
+    `False` and pixels without clouds are `True`.
     '''
     with TemporaryDirectory() as tmp_dir:
         # we need the temporary directory to work around a problem with reading
@@ -137,10 +145,25 @@ def scihub_cloud_mask(product_path):
                     warnings.simplefilter("ignore")
                     # this returns a warning because the iterator has to be
                     # rewound; while this is a performance issue, we can ignore it
-                    return unary_union([shape(f['geometry']) for f in features])
+                    mask = unary_union([shape(f['geometry']) for f in features])
         except ValueError:
             # empty cloud mask
-            return Polygon([])
+            mask = Polygon([])
+            
+    if kwargs.get('rasterize'):
+        # return raster version of the vector geometry we found above
+        target_shape = kwargs.get('target_shape')
+        target_transform = kwargs.get('target_transform')
+        if not target_transform or not target_shape:
+            error_msg = 'target_transform and target_shape need to be set ' + \
+                'to construct a rasterized cloud mask.'
+            raise ValueError(error_msg)
+        
+        return geometry_mask(mask,
+                             out_shape=target_shape,
+                             transform=target_transform)
+    else:
+        return mask
     
 
 def scihub_normalize_range(v):