"""Common tasks shared between pipelines."""
from __future__ import annotations
from functools import partial
from typing import Any
import cv2
import numpy as np
from eolearn.core import EOPatch, EOTask, MapFeatureTask, SaveTask
from eolearn.core.types import Feature, FeaturesSpecification
from eolearn.core.utils.parsing import parse_renamed_feature
from eolearn.geometry import MorphologicalOperations
[docs]class ClassFilterTask(EOTask):
"""Run class specific morphological operation."""
def __init__(
self,
feature: Feature,
labels: list[int],
morph_operation: MorphologicalOperations,
struct_elem: np.ndarray | None = None,
):
"""Perform a morphological operation on a given feature mask
:param feature: Feature to be modified
:param labels: List of labels to be considered for morphological operation
:param morph_operation: Type of morphological operation ot perform
:param struct_elem: Structured element to be used. Taken from `ml_tools.MorphologicalStructFactory`
"""
self.feature_name: str | None
self.new_feature_name: str | None
self.renamed_feature = parse_renamed_feature(feature)
self.labels = labels
self.morph_operation = MorphologicalOperations.get_operation(morph_operation)
self.struct_elem = struct_elem
[docs] def execute(self, eopatch: EOPatch) -> EOPatch:
feature_type, feature_name, new_feature_name = self.renamed_feature
mask = eopatch[(feature_type, feature_name)].copy()
morp_func = partial(cv2.morphologyEx, kernel=self.struct_elem, op=self.morph_operation)
for label in self.labels:
label_mask = np.squeeze((mask == label).astype(np.uint8), axis=-1)
mask_mod: np.ndarray = morp_func(label_mask) * label # type: ignore[assignment]
mask_mod = mask_mod[..., np.newaxis]
mask[mask == label] = mask_mod[mask == label]
eopatch[(feature_type, new_feature_name)] = mask
return eopatch
[docs]class LinearFunctionTask(MapFeatureTask):
"""Applies a linear function to the values of input features.
Each value in the feature is modified as `x -> x * slope + intercept`. The `dtype` of the result can be customized.
"""
def __init__(
self,
input_features: FeaturesSpecification,
output_features: FeaturesSpecification | None = None,
slope: float = 1,
intercept: float = 0,
dtype: str | type | np.dtype | None = None,
):
"""
:param input_features: Feature or features on which the function is used.
:param output_features: Feature or features for saving the result. If not provided the input_features are
overwritten.
:param slope: Slope of the function i.e. the multiplication factor.
:param intercept: Intercept of the function i.e. the value added.
:param dtype: Numpy dtype of the output feature. If not provided the dtype is determined by Numpy, so it is
recommended to set manually.
"""
if output_features is None:
output_features = input_features
self.dtype = dtype if dtype is None else np.dtype(dtype)
super().__init__(input_features, output_features, slope=slope, intercept=intercept)
[docs] def map_method(self, feature: np.ndarray, slope: float, intercept: float) -> np.ndarray: # type:ignore[override]
"""A method where feature is multiplied by a slope"""
rescaled_feature = feature * slope + intercept
return rescaled_feature if self.dtype is None else rescaled_feature.astype(self.dtype)
[docs]class SkippableSaveTask(SaveTask):
"""Same as `SaveTask` but can be skipped if the `eopatch_folder` is set to `None`."""
[docs] def execute(self, eopatch: EOPatch, *, eopatch_folder: str | None = "", **kwargs: Any) -> EOPatch:
if eopatch_folder is None:
return eopatch
return super().execute(eopatch, eopatch_folder=eopatch_folder, **kwargs)