handle "classification" bands (#428)

geetools/ee_image.py

@@ -21,7 +21,7 @@
 from xee.ext import REQUEST_BYTE_LIMIT
 
 from .accessors import register_class_accessor
-from .utils import plot_data
+from .utils import format_class_info, plot_data
 
 
 @register_class_accessor(ee.Image, "geetools")
@@ -1678,6 +1678,76 @@ def fromList(images: ee.List | list[ee.Image]) -> ee.Image:
         ic = ee.ImageCollection.fromImages(images)
         return ic.toBands().rename(bandNames)
 
+    def classToBands(self, class_info: dict, band: str | int = 0) -> ee.Image:
+        """Convert each class into a separate binary mask.
+
+        Args:
+            class_info: class information.
+            band: band that contains class information
+
+        Example:
+            .. jupyter-execute::
+
+                import ee, geetools
+                from geetools.utils import initialize_documentation
+
+                initialize_documentation()
+
+                class_info = {
+                  '2': 'dark',
+                  '3': 'shadow',
+                  '7': 'clouds_low',
+                  '8': 'clouds_mid',
+                  '9': 'clouds_high',
+                  '10': 'cirrus'
+                }
+                image = ee.Image("COPERNICUS/S2_SR_HARMONIZED/20230120T142709_20230120T143451_T18GYT")
+                decoded = image.geetools.classToBands(class_info, "SCL")
+                decoded.getInfo()
+        """
+        class_info = format_class_info(class_info)
+        # I don't see the class info coming from the server, so it'll client side until I get challenged
+        images = []
+        for class_value, class_name in class_info.items():
+            class_value = int(class_value)
+            mask = self._obj.select([band]).eq(class_value).rename(class_name)
+            images.append(mask)
+        return self.fromList(images)
+
+    def classMask(self, class_info: dict, classes: Optional[list] = None, band: str | int = 0):
+        """Create a mask using the class information.
+
+         In this case there is no option for "all" or "any" because class bands cannot contain 2 classes in the same pixel value.
+
+        Args:
+            class_info: class information (client-side only).
+            classes: name of the classes to use for the mask. If None it will use all classes in bits_info.
+            band: name of the bit band. Defaults to first band.
+
+        Example:
+            .. jupyter-execute::
+
+                import ee, geetools
+                from geetools.utils import initialize_documentation
+
+                initialize_documentation()
+
+                class_info = {
+                  '2': 'dark',
+                  '3': 'shadow',
+                  '7': 'clouds_low',
+                  '8': 'clouds_mid',
+                  '9': 'clouds_high',
+                  '10': 'cirrus'
+                }
+                image = ee.Image("COPERNICUS/S2_SR_HARMONIZED/20230120T142709_20230120T143451_T18GYT")
+                decoded = image.geetools.classMask(class_info, classes=["dark", "shadow"], band="SCL")
+                decoded.getInfo()
+        """
+        masks = self.classToBands(class_info, band)
+        masks = masks.select(classes) if classes else masks
+        return masks.reduce(ee.Reducer.anyNonZero()).rename("mask")
+
     def byBands(
         self,
         regions: ee.FeatureCollection,

geetools/utils.py

@@ -267,3 +267,22 @@ def initialize_documentation():
         )
 
     pass
+
+
+def format_bandname(name: str, replacement: str = "_") -> str:
+    """Format a band name to be allowed in GEE."""
+    banned = list(".*/¿?[]{}+#$%&")
+    return str([name.replace(char, replacement) for char in banned][0])
+
+
+def format_class_info(class_info: dict) -> dict:
+    """Format the class information.
+
+    Args:
+        class_info: class information in a dict ({class value: class name})
+    """
+    final = {}
+    for class_value, class_name in class_info.items():
+        # make sure class value is an int, but store as str
+        final[str(int(class_value))] = format_bandname(class_name)
+    return final

tests/conftest.py

@@ -276,3 +276,27 @@ def stac_schema():
 def jaxa_rainfall():
     """Return the JAXA rain collection."""
     return ee.ImageCollection("JAXA/GPM_L3/GSMaP/v6/operational")
+
+
+@pytest.fixture
+def s2_class_image():
+    """A Sentinel 2 image that contains clouds, shadows, water and snow.
+
+    This image is located in Argentina and Chile (Patagonia).
+    """
+    return ee.Image("COPERNICUS/S2_SR_HARMONIZED/20230120T142709_20230120T143451_T18GYT")
+
+
+@pytest.fixture
+def polygon_instance():
+    """Return a defined polygon instance."""
+    return ee.Geometry.Polygon(
+        [
+            [
+                [-71.84476689765019, -42.81816243454466],
+                [-71.84476689765019, -42.897690198549135],
+                [-71.72391728827519, -42.897690198549135],
+                [-71.72391728827519, -42.81816243454466],
+            ]
+        ]
+    )

tests/test_Image.py

@@ -618,6 +618,33 @@ def image(self):
         return image.select(["B4", "B3", "B2"])
 
 
+class TestClassToBands:
+    """Test the ``classToBands`` method."""
+
+    def test_class_to_bands(self, s2_class_image, polygon_instance, ee_image_regression):
+        """Test the ``classToBands`` method."""
+        class_info = {"3": "shadow", "9": "clouds_high", "11": "snow"}
+        decoded = s2_class_image.geetools.classToBands(class_info, "SCL")
+        ee_image_regression.check(decoded, scale=10, region=polygon_instance)
+
+
+class TestClassMask:
+    """Test the ``classMask`` method."""
+
+    def test_class_mask(self, s2_class_image, polygon_instance, ee_image_regression):
+        """Test the ``classMask`` method."""
+        class_info = {
+            "2": "dark",
+            "3": "shadow",
+            "7": "clouds_low",
+            "8": "clouds_mid",
+            "9": "clouds_high",
+            "10": "cirrus",
+        }
+        decoded = s2_class_image.geetools.classMask(class_info, band="SCL")
+        ee_image_regression.check(decoded, scale=10, region=polygon_instance)
+
+
 class TestPlot:
     """Test the ``plot`` method."""
 

tests/test_Image/serialized_test_class_mask.yml

@@ -0,0 +1,200 @@
+result: '0'
+values:
+  '0':
+    functionInvocationValue:
+      arguments:
+        geometry:
+          functionInvocationValue:
+            arguments:
+              coordinates:
+                constantValue:
+                - - - -71.84476689765019
+                    - -42.81816243454466
+                  - - -71.84476689765019
+                    - -42.897690198549135
+                  - - -71.72391728827519
+                    - -42.897690198549135
+                  - - -71.72391728827519
+                    - -42.81816243454466
+              evenOdd:
+                constantValue: true
+            functionName: GeometryConstructors.Polygon
+        input:
+          functionInvocationValue:
+            arguments:
+              input:
+                functionInvocationValue:
+                  arguments:
+                    image:
+                      functionInvocationValue:
+                        arguments:
+                          input:
+                            functionInvocationValue:
+                              arguments:
+                                collection:
+                                  functionInvocationValue:
+                                    arguments:
+                                      images:
+                                        valueReference: '1'
+                                    functionName: ImageCollection.fromImages
+                              functionName: ImageCollection.toBands
+                          names:
+                            functionInvocationValue:
+                              arguments:
+                                list:
+                                  functionInvocationValue:
+                                    arguments:
+                                      baseAlgorithm:
+                                        functionDefinitionValue:
+                                          argumentNames:
+                                          - _MAPPING_VAR_0_0
+                                          body: '3'
+                                      dropNulls:
+                                        constantValue: false
+                                      list:
+                                        valueReference: '1'
+                                    functionName: List.map
+                              functionName: List.flatten
+                        functionName: Image.rename
+                    reducer:
+                      functionInvocationValue:
+                        arguments: {}
+                        functionName: Reducer.anyNonZero
+                  functionName: Image.reduce
+              names:
+                constantValue:
+                - mask
+            functionName: Image.rename
+        scale:
+          constantValue: 10
+      functionName: Image.clipToBoundsAndScale
+  '1':
+    arrayValue:
+      values:
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 2
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - dark
+          functionName: Image.rename
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 3
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - shadow
+          functionName: Image.rename
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 7
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - clouds_low
+          functionName: Image.rename
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 8
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - clouds_mid
+          functionName: Image.rename
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 9
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - clouds_high
+          functionName: Image.rename
+      - functionInvocationValue:
+          arguments:
+            input:
+              functionInvocationValue:
+                arguments:
+                  image1:
+                    valueReference: '2'
+                  image2:
+                    functionInvocationValue:
+                      arguments:
+                        value:
+                          constantValue: 10
+                      functionName: Image.constant
+                functionName: Image.eq
+            names:
+              constantValue:
+              - cirrus
+          functionName: Image.rename
+  '2':
+    functionInvocationValue:
+      arguments:
+        bandSelectors:
+          constantValue:
+          - SCL
+        input:
+          functionInvocationValue:
+            arguments:
+              id:
+                constantValue: COPERNICUS/S2_SR_HARMONIZED/20230120T142709_20230120T143451_T18GYT
+            functionName: Image.load
+      functionName: Image.select
+  '3':
+    functionInvocationValue:
+      arguments:
+        image:
+          argumentReference: _MAPPING_VAR_0_0
+      functionName: Image.bandNames