Merge pull request #35 from digitalearthpacific/implement-filters-s2

Move S2Processor to Processors and add filter to s2masking

dep_tools/processors.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Tuple
 
 from xarray import DataArray
 
-from .landsat_utils import mask_clouds
+from .landsat_utils import mask_clouds as mask_clouds_landsat
+from .s2_utils import harmonize_to_old
+from .s2_utils import mask_clouds as mask_clouds_s2
 from .utils import scale_and_offset
 
 
@@ -31,7 +32,7 @@ def __init__(
 
     def process(self, xr: DataArray) -> DataArray:
         if self.mask_clouds:
-            xr = mask_clouds(xr, **self.mask_kwargs)
+            xr = mask_clouds_landsat(xr, **self.mask_kwargs)
         if self.scale_and_offset:
             # These values only work for SR bands of landsat. Ideally we could
             # read from metadata. _Really_ ideally we could just pass "scale"
@@ -41,3 +42,26 @@ def process(self, xr: DataArray) -> DataArray:
             xr = scale_and_offset(xr, scale=[scale], offset=offset)
 
         return xr
+
+
+class S2Processor(Processor):
+    def __init__(
+        self,
+        send_area_to_processor: bool = False,
+        scale_and_offset: bool = True,
+        mask_clouds: bool = True,
+        mask_clouds_kwargs: dict = dict(),
+    ) -> None:
+        super().__init__(send_area_to_processor)
+        self.scale_and_offset = scale_and_offset
+        self.mask_clouds = mask_clouds
+        self.mask_kwargs = mask_clouds_kwargs
+
+    def process(self, xr: DataArray) -> DataArray:
+        if self.mask_clouds:
+            xr = mask_clouds_s2(xr, **self.mask_kwargs)
+
+        if self.scale_and_offset:
+            xr = harmonize_to_old(xr)
+
+        return xr

dep_tools/s2_utils.py

@@ -1,43 +1,14 @@
 import datetime
+from typing import Iterable, Tuple
 
 from odc.algo import erase_bad, mask_cleanup
 from xarray import DataArray, concat
 
-from .processors import Processor
-
-
-class S2Processor(Processor):
-    def __init__(
-        self,
-        send_area_to_processor: bool = False,
-        scale_and_offset: bool = True,
-        mask_clouds: bool = True,
-        dilate_mask: bool = False,
-        keep_ints: bool = False,
-    ) -> None:
-        super().__init__(send_area_to_processor)
-        self.scale_and_offset = scale_and_offset
-        self.mask_clouds = mask_clouds
-        self.dilate_mask = dilate_mask
-        self.keep_ints = keep_ints
-
-    def process(self, xr: DataArray) -> DataArray:
-        if self.mask_clouds:
-            xr = mask_clouds(xr, dilate=self.dilate_mask, keep_ints=self.keep_ints)
-
-        if self.scale_and_offset:
-            xr = scale_and_offset_s2(xr)
-
-        return xr
-
-
-def scale_and_offset_s2(da: DataArray) -> DataArray:
-    # TODO: don't scale SCL
-    return harmonize_to_old(da) * 0.0001
-
 
 def mask_clouds(
-    xr: DataArray, dilate: bool = True, keep_ints: bool = False
+    xr: DataArray,
+    filters: Iterable[Tuple[str, int]] | None = None,
+    keep_ints: bool = False,
 ) -> DataArray:
     # NO_DATA = 0
     # SATURATED_OR_DEFECTIVE = 1
@@ -61,9 +32,8 @@ def mask_clouds(
     except KeyError:
         cloud_mask = xr.SCL.astype("uint16") & bitmask != 0
 
-    if dilate:
-        # From Alex @ https://gist.github.com/alexgleith/d9ea655d4e55162e64fe2c9db84284e5
-        cloud_mask = mask_cleanup(cloud_mask, [("opening", 2), ("dilation", 3)])
+    if filters is not None:
+        cloud_mask = mask_cleanup(cloud_mask, filters)
 
     if keep_ints:
         return erase_bad(xr, cloud_mask)