Merge pull request #91 from WorldCereal/86-generete-map-params

kvantricht · web-flow · commit fefd73d8991d · 2024-07-02T12:17:58.000+02:00
Added configuration data classes for cropland/croptype mapping workflow
diff --git a/scripts/inference/cropland_mapping.py b/scripts/inference/cropland_mapping.py
@@ -20,25 +20,25 @@
     parser.add_argument("maxx", type=float, help="Maximum X coordinate (east)")
     parser.add_argument("maxy", type=float, help="Maximum Y coordinate (north)")
     parser.add_argument(
-        "--epsg",
-        type=int,
-        default=4326,
-        help="EPSG code of the input `minx`, `miny`, `maxx`, `maxy` parameters.",
+        "start_date", type=str, help="Starting date for data extraction."
     )
+    parser.add_argument("end_date", type=str, help="Ending date for data extraction.")
     parser.add_argument(
         "product",
         type=str,
         help="Product to generate. One of ['cropland', 'croptype']",
     )
-    parser.add_argument(
-        "start_date", type=str, help="Starting date for data extraction."
-    )
-    parser.add_argument("end_date", type=str, help="Ending date for data extraction.")
     parser.add_argument(
         "output_path",
         type=Path,
         help="Path to folder where to save the resulting GeoTiff.",
     )
+    parser.add_argument(
+        "--epsg",
+        type=int,
+        default=4326,
+        help="EPSG code of the input `minx`, `miny`, `maxx`, `maxy` parameters.",
+    )
 
     args = parser.parse_args()
 
diff --git a/src/worldcereal/job.py b/src/worldcereal/job.py
@@ -1,19 +1,16 @@
 """Executing inference jobs on the OpenEO backend."""
 
-from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 from typing import Optional, Union
 
 import openeo
-from openeo import DataCube
 from openeo_gfmap import BackendContext, BoundingBoxExtent, TemporalContext
-from openeo_gfmap.features.feature_extractor import apply_feature_extractor
-from openeo_gfmap.inference.model_inference import apply_model_inference
-from openeo_gfmap.preprocessing.scaling import compress_uint8, compress_uint16
+from pydantic import BaseModel
 
 from worldcereal.openeo.feature_extractor import PrestoFeatureExtractor
 from worldcereal.openeo.inference import CroplandClassifier, CroptypeClassifier
+from worldcereal.openeo.mapping import _cropland_map, _croptype_map
 from worldcereal.openeo.preprocessing import worldcereal_preprocessed_inputs_gfmap
 
 
@@ -24,44 +21,109 @@ class WorldCerealProduct(Enum):
     CROPTYPE = "croptype"
 
 
-ONNX_DEPS_URL = "https://artifactory.vgt.vito.be/artifactory/auxdata-public/openeo/onnx_dependencies_1.16.3.zip"
+class FeaturesParameters(BaseModel):
+    """Parameters for the feature extraction UDFs. Types are enforced by
+    Pydantic.
 
-PRODUCT_SETTINGS = {
-    WorldCerealProduct.CROPLAND: {
-        "features": {
-            "extractor": PrestoFeatureExtractor,
-            "parameters": {
-                "rescale_s1": False,  # Will be done in the Presto UDF itself!
-                "presto_model_url": "https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal-minimal-inference/presto.pt",  # NOQA
-            },
-        },
-        "classification": {
-            "classifier": CroplandClassifier,
-            "parameters": {
-                "classifier_url": "https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal-minimal-inference/wc_catboost.onnx"  # NOQA
-            },
-        },
-    },
-    WorldCerealProduct.CROPTYPE: {
-        "features": {
-            "extractor": PrestoFeatureExtractor,
-            "parameters": {
-                "rescale_s1": False,  # Will be done in the Presto UDF itself!
-                "presto_model_url": "https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal/models/PhaseII/presto-ss-wc-ft-ct-30D_test.pt",  # NOQA
-            },
-        },
-        "classification": {
-            "classifier": CroptypeClassifier,
-            "parameters": {
-                "classifier_url": "https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal/models/PhaseII/presto-ss-wc-ft-ct-30D_test_CROPTYPE9.onnx"  # NOQA
-            },
-        },
-    },
-}
+    Attributes
+    ----------
+    rescale_s1 : bool (default=False)
+        Whether to rescale Sentinel-1 bands before feature extraction. Should be
+        left to False, as this is done in the Presto UDF itself.
+    presto_model_url : str
+        Public URL to the Presto model used for feature extraction. The file
+        should be a PyTorch serialized model.
+    """
+
+    rescale_s1: bool
+    presto_model_url: str
 
 
-@dataclass
-class InferenceResults:
+class ClassifierParameters(BaseModel):
+    """Parameters for the classifier. Types are enforced by Pydantic.
+
+    Attributes
+    ----------
+    classifier_url : str
+        Public URL to the classifier model. Te file should be an ONNX accepting
+        a `features` field for input data and returning either two output
+        probability arrays `true` and `false` in case of cropland mapping, or
+        a probability array per-class in case of croptype mapping.
+    """
+
+    classifier_url: str
+
+
+class CropLandParameters(BaseModel):
+    """Parameters for the cropland product inference pipeline. Types are
+    enforced by Pydantic.
+
+    Attributes
+    ----------
+    feature_extractor : PrestoFeatureExtractor
+        Feature extractor to use for the inference. This class must be a
+        subclass of GFMAP's `PatchFeatureExtractor` and returns float32
+        features.
+    features_parameters : FeaturesParameters
+        Parameters for the feature extraction UDF. Will be serialized into a
+        dictionnary and passed in the process graph.
+    classifier : CroplandClassifier
+        Classifier to use for the inference. This class must be a subclass of
+        GFMAP's `ModelInference` and returns predictions/probabilities for
+        cropland.
+    classifier_parameters : ClassifierParameters
+        Parameters for the classifier UDF. Will be serialized into a dictionnary
+        and passed in the process graph.
+    """
+
+    feature_extractor: PrestoFeatureExtractor = PrestoFeatureExtractor
+    features_parameters: FeaturesParameters = FeaturesParameters(
+        rescale_s1=False,
+        presto_model_url="https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal-minimal-inference/presto.pt",  # NOQA
+    )
+    classifier: CroplandClassifier = CroplandClassifier
+    classifier_parameters: ClassifierParameters = ClassifierParameters(
+        classifier_url="https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal-minimal-inference/wc_catboost.onnx"  # NOQA
+    )
+
+
+class CropTypeParameters(BaseModel):
+    """Parameters for the croptype product inference pipeline. Types are
+    enforced by Pydantic.
+
+    Attributes
+    ----------
+    feature_extractor : PrestoFeatureExtractor
+        Feature extractor to use for the inference. This class must be a
+        subclass of GFMAP's `PatchFeatureExtractor` and returns float32
+        features.
+    features_parameters : FeaturesParameters
+        Parameters for the feature extraction UDF. Will be serialized into a
+        dictionnary and passed in the process graph.
+    classifier : CroptypeClassifier
+        Classifier to use for the inference. This class must be a subclass of
+        GFMAP's `ModelInference` and returns predictions/probabilities for
+        cropland classes.
+    classifier_parameters : ClassifierParameters
+        Parameters for the classifier UDF. Will be serialized into a dictionnary
+        and passed in the process graph.
+    """
+
+    feature_extractor: PrestoFeatureExtractor = PrestoFeatureExtractor
+    feature_parameters: FeaturesParameters = FeaturesParameters(
+        rescale_s1=False,
+        presto_model_url="https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal/models/PhaseII/presto-ss-wc-ft-ct-30D_test.pt",  # NOQA
+    )
+    classifier: CroptypeClassifier = CroptypeClassifier
+    classifier_parameters: ClassifierParameters = ClassifierParameters(
+        classifier_url="https://artifactory.vgt.vito.be/artifactory/auxdata-public/worldcereal/models/PhaseII/presto-ss-wc-ft-ct-30D_test_CROPTYPE9.onnx"  # NOQA
+    )
+
+
+ONNX_DEPS_URL = "https://artifactory.vgt.vito.be/artifactory/auxdata-public/openeo/onnx_dependencies_1.16.3.zip"
+
+
+class InferenceResults(BaseModel):
     """Dataclass to store the results of the WorldCereal job.
 
     Attributes
@@ -88,6 +150,8 @@ def generate_map(
     backend_context: BackendContext,
     output_path: Optional[Union[Path, str]],
     product_type: WorldCerealProduct = WorldCerealProduct.CROPLAND,
+    cropland_parameters: CropLandParameters = CropLandParameters(),
+    croptype_parameters: Optional[CropTypeParameters] = CropTypeParameters(),
     out_format: str = "GTiff",
 ) -> InferenceResults:
     """Main function to generate a WorldCereal product.
@@ -104,6 +168,12 @@ def generate_map(
         output path to download the product to
     product_type : WorldCerealProduct, optional
         product describer, by default WorldCerealProduct.CROPLAND
+    cropland_parameters: CropLandParameters
+        Parameters for the cropland product inference pipeline.
+    croptype_parameters: Optional[CropTypeParameters]
+        Parameters for the croptype product inference pipeline. Only required
+        whenever `product_type` is set to `WorldCerealProduct.CROPTYPE`, will be
+        ignored otherwise.
     out_format : str, optional
         Output format, by default "GTiff"
 
@@ -120,7 +190,7 @@ def generate_map(
         if the out_format is not supported
     """
 
-    if product_type not in PRODUCT_SETTINGS.keys():
+    if product_type not in WorldCerealProduct:
         raise ValueError(f"Product {product_type.value} not supported.")
 
     if out_format not in ["GTiff", "NetCDF"]:
@@ -145,18 +215,22 @@ def generate_map(
 
     # Construct the feature extraction and model inference pipeline
     if product_type == WorldCerealProduct.CROPLAND:
-        classes = _cropland_map(inputs)
+        classes = _cropland_map(inputs, cropland_parameters=cropland_parameters)
     elif product_type == WorldCerealProduct.CROPTYPE:
         # First compute cropland map
         cropland_mask = (
-            _cropland_map(inputs)
+            _cropland_map(inputs, cropland_parameters=cropland_parameters)
             .filter_bands("classification")
             .reduce_dimension(
                 dimension="t", reducer="mean"
             )  # Temporary fix to make this work as mask
         )
 
-        classes = _croptype_map(inputs, cropland_mask=cropland_mask)
+        classes = _croptype_map(
+            inputs,
+            croptype_parameters=croptype_parameters,
+            cropland_mask=cropland_mask,
+        )
 
     # Submit the job
     job = classes.execute_batch(
@@ -219,129 +293,3 @@ def collect_inputs(
         out_format="NetCDF",
         job_options={"driver-memory": "4g", "executor-memoryOverhead": "4g"},
     )
-
-
-def _cropland_map(inputs: DataCube) -> DataCube:
-    """Method to produce cropland map from preprocessed inputs, using
-    a Presto feature extractor and a CatBoost classifier.
-
-    Parameters
-    ----------
-    inputs : DataCube
-        preprocessed input cube
-
-    Returns
-    -------
-    DataCube
-        binary labels and probability
-    """
-
-    # Run feature computer
-    features = apply_feature_extractor(
-        feature_extractor_class=PRODUCT_SETTINGS[WorldCerealProduct.CROPLAND][
-            "features"
-        ]["extractor"],
-        cube=inputs,
-        parameters=PRODUCT_SETTINGS[WorldCerealProduct.CROPLAND]["features"][
-            "parameters"
-        ],
-        size=[
-            {"dimension": "x", "unit": "px", "value": 100},
-            {"dimension": "y", "unit": "px", "value": 100},
-        ],
-        overlap=[
-            {"dimension": "x", "unit": "px", "value": 0},
-            {"dimension": "y", "unit": "px", "value": 0},
-        ],
-    )
-
-    # Run model inference on features
-    classes = apply_model_inference(
-        model_inference_class=PRODUCT_SETTINGS[WorldCerealProduct.CROPLAND][
-            "classification"
-        ]["classifier"],
-        cube=features,
-        parameters=PRODUCT_SETTINGS[WorldCerealProduct.CROPLAND]["classification"][
-            "parameters"
-        ],
-        size=[
-            {"dimension": "x", "unit": "px", "value": 100},
-            {"dimension": "y", "unit": "px", "value": 100},
-            {"dimension": "t", "value": "P1D"},
-        ],
-        overlap=[
-            {"dimension": "x", "unit": "px", "value": 0},
-            {"dimension": "y", "unit": "px", "value": 0},
-        ],
-    )
-
-    # Cast to uint8
-    classes = compress_uint8(classes)
-
-    return classes
-
-
-def _croptype_map(inputs: DataCube, cropland_mask: DataCube = None) -> DataCube:
-    """Method to produce croptype map from preprocessed inputs, using
-    a Presto feature extractor and a CatBoost classifier.
-
-    Parameters
-    ----------
-    inputs : DataCube
-        preprocessed input cube
-    cropland_mask : DataCube, optional
-        optional cropland mask, by default None
-
-    Returns
-    -------
-    DataCube
-        croptype labels and probability
-    """
-
-    # Run feature computer
-    features = apply_feature_extractor(
-        feature_extractor_class=PRODUCT_SETTINGS[WorldCerealProduct.CROPTYPE][
-            "features"
-        ]["extractor"],
-        cube=inputs,
-        parameters=PRODUCT_SETTINGS[WorldCerealProduct.CROPTYPE]["features"][
-            "parameters"
-        ],
-        size=[
-            {"dimension": "x", "unit": "px", "value": 100},
-            {"dimension": "y", "unit": "px", "value": 100},
-        ],
-        overlap=[
-            {"dimension": "x", "unit": "px", "value": 0},
-            {"dimension": "y", "unit": "px", "value": 0},
-        ],
-    )
-
-    # Run model inference on features
-    classes = apply_model_inference(
-        model_inference_class=PRODUCT_SETTINGS[WorldCerealProduct.CROPTYPE][
-            "classification"
-        ]["classifier"],
-        cube=features,
-        parameters=PRODUCT_SETTINGS[WorldCerealProduct.CROPTYPE]["classification"][
-            "parameters"
-        ],
-        size=[
-            {"dimension": "x", "unit": "px", "value": 100},
-            {"dimension": "y", "unit": "px", "value": 100},
-            {"dimension": "t", "value": "P1D"},
-        ],
-        overlap=[
-            {"dimension": "x", "unit": "px", "value": 0},
-            {"dimension": "y", "unit": "px", "value": 0},
-        ],
-    )
-
-    # Mask cropland
-    if cropland_mask is not None:
-        classes = classes.mask(cropland_mask == 0, replacement=0)
-
-    # Cast to uint16
-    classes = compress_uint16(classes)
-
-    return classes
diff --git a/src/worldcereal/openeo/mapping.py b/src/worldcereal/openeo/mapping.py
