Merge pull request #438 from sentinel-hub/feat/simple-filter-task

Updated SimpleFilterTask and wrote additional unit tests

Author: zigaLuksic

features/eolearn/features/feature_manipulation.py

@@ -15,11 +15,12 @@
 import datetime as dt
 import logging
 from functools import partial
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
+from geopandas import GeoDataFrame
 
-from eolearn.core import EOPatch, EOTask, FeatureType, MapFeatureTask
+from eolearn.core import EOPatch, EOTask, FeatureType, FeatureTypeSet, MapFeatureTask
 
 from .utils import ResizeLib, ResizeMethod, spatially_resize_image
 
@@ -34,51 +35,57 @@ class SimpleFilterTask(EOTask):
     A filter_func is a callable which takes a numpy array and returns a bool.
     """
 
-    def __init__(self, feature, filter_func, filter_features=...):
+    def __init__(
+        self,
+        feature: Union[Tuple[FeatureType, str], FeatureType],
+        filter_func: Union[Callable[[np.ndarray], bool], Callable[[dt.datetime], bool]],
+        filter_features: Any = ...,
+    ):
         """
         :param feature: Feature in the EOPatch , e.g. feature=(FeatureType.DATA, 'bands')
-        :type feature: (FeatureType, str)
         :param filter_func: A callable that takes a numpy evaluates to bool.
-        :type filter_func: object
-        :param filter_features: A collection of features which will be filtered
-        :type filter_features: dict(FeatureType: set(str))
+        :param filter_features: A collection of features which will be filtered into a new EOPatch
         """
-        self.feature = self.parse_feature(feature)
+        self.feature = self.parse_feature(
+            feature, allowed_feature_types=FeatureTypeSet.TEMPORAL_TYPES.difference([FeatureType.VECTOR])
+        )
         self.filter_func = filter_func
         self.filter_features_parser = self.get_feature_parser(filter_features)
 
-    def _get_filtered_indices(self, feature_data):
+    def _get_filtered_indices(self, feature_data: Iterable) -> List[int]:
+        """Get valid time indices from either a numpy array or a list of timestamps."""
         return [idx for idx, img in enumerate(feature_data) if self.filter_func(img)]
 
-    def _update_other_data(self, eopatch):
-        pass
+    @staticmethod
+    def _filter_vector_feature(gdf: GeoDataFrame, good_idxs: List[int], timestamps: List[dt.datetime]) -> GeoDataFrame:
+        """Filters rows that don't match with the timestamps that will be kept."""
+        timestamps_to_keep = set(timestamps[idx] for idx in good_idxs)
+        return gdf[gdf.TIMESTAMP.isin(timestamps_to_keep)]
 
-    def execute(self, eopatch):
+    def execute(self, eopatch: EOPatch) -> EOPatch:
         """
-        :param eopatch: Input EOPatch.
-        :type eopatch: EOPatch
-        :return: Transformed eo patch
-        :rtype: EOPatch
+        :param eopatch: An input EOPatch.
+        :return: A new EOPatch with filtered features.
         """
+        filtered_eopatch = EOPatch()
         good_idxs = self._get_filtered_indices(eopatch[self.feature])
-        if not good_idxs:
-            raise RuntimeError("EOPatch has no good indices after filtering with given filter function")
-
-        for feature_type, feature_name in self.filter_features_parser.get_features(eopatch):
-            if feature_type.is_temporal():
-                if feature_type.has_dict():
-                    if feature_type.contains_ndarrays():
-                        eopatch[feature_type][feature_name] = np.asarray(
-                            [eopatch[feature_type][feature_name][idx] for idx in good_idxs]
-                        )
-                    # else:
-                    #     NotImplemented
+
+        for feature in self.filter_features_parser.get_features(eopatch):
+            feature_type, _ = feature
+            data = eopatch[feature]
+
+            if feature_type is FeatureType.TIMESTAMP:
+                data = [data[idx] for idx in good_idxs]
+
+            elif feature_type.is_temporal():
+                if feature_type.is_raster():
+                    data = data[good_idxs]
                 else:
-                    eopatch[feature_type] = [eopatch[feature_type][idx] for idx in good_idxs]
+                    data = self._filter_vector_feature(data, good_idxs, eopatch.timestamp)
 
-        self._update_other_data(eopatch)
+            filtered_eopatch[feature] = data
 
-        return eopatch
+        return filtered_eopatch
 
 
 class FilterTimeSeriesTask(SimpleFilterTask):

features/eolearn/tests/test_feature_manipulation.py

@@ -16,11 +16,39 @@
 from numpy.testing import assert_allclose, assert_array_equal
 
 from eolearn.core import EOPatch, FeatureType
-from eolearn.features import FilterTimeSeriesTask, LinearFunctionTask, ValueFilloutTask
+from eolearn.features import FilterTimeSeriesTask, LinearFunctionTask, SimpleFilterTask, ValueFilloutTask
 from eolearn.features.feature_manipulation import SpatialResizeTask
 
 
-def test_content_after_timefilter():
+@pytest.mark.parametrize(
+    "feature", [(FeatureType.DATA, "BANDS-S2-L1C"), FeatureType.TIMESTAMP, (FeatureType.LABEL, "IS_CLOUDLESS")]
+)
+def test_simple_filter_task_filter_all(example_eopatch: EOPatch, feature):
+    filter_all_task = SimpleFilterTask(feature, filter_func=lambda _: False)
+    filtered_eopatch = filter_all_task.execute(example_eopatch)
+
+    assert filtered_eopatch is not example_eopatch
+    assert filtered_eopatch.data["CLP"].shape == (0, 101, 100, 1)
+    assert filtered_eopatch.scalar["CLOUD_COVERAGE"].shape == (0, 1)
+    assert len(filtered_eopatch.vector["CLM_VECTOR"]) == 0
+    assert np.array_equal(filtered_eopatch.mask_timeless["LULC"], example_eopatch.mask_timeless["LULC"])
+    assert filtered_eopatch.timestamp == []
+
+
+@pytest.mark.parametrize(
+    "feature", [(FeatureType.MASK, "CLM"), FeatureType.TIMESTAMP, (FeatureType.SCALAR, "CLOUD_COVERAGE")]
+)
+def test_simple_filter_task_filter_nothing(example_eopatch: EOPatch, feature):
+    del example_eopatch.data["REFERENCE_SCENES"]  # Wrong size of time dimension
+
+    filter_all_task = SimpleFilterTask(feature, filter_func=lambda _: True)
+    filtered_eopatch = filter_all_task.execute(example_eopatch)
+
+    assert filtered_eopatch is not example_eopatch
+    assert filtered_eopatch == example_eopatch
+
+
+def test_content_after_time_filter():
     timestamps = [
         datetime.datetime(2017, 1, 1, 10, 4, 7),
         datetime.datetime(2017, 1, 4, 10, 14, 5),
@@ -37,16 +65,14 @@ def test_content_after_timefilter():
 
     new_start, new_end = 4, -3
 
-    new_interval = (timestamps[new_start], timestamps[new_end])
-
-    new_timestamps = timestamps[new_start : new_end + 1]
-
     eop = EOPatch(timestamp=timestamps, data={"data": data})
 
-    filter_task = FilterTimeSeriesTask(start_date=new_interval[0], end_date=new_interval[1])
-    filter_task.execute(eop)
+    filter_task = FilterTimeSeriesTask(start_date=timestamps[new_start], end_date=timestamps[new_end])
+    filtered_eop = filter_task.execute(eop)
 
-    assert new_timestamps == eop.timestamp
+    assert filtered_eop is not eop
+    assert filtered_eop.timestamp == timestamps[new_start : new_end + 1]
+    assert np.array_equal(filtered_eop.data["data"], data[new_start : new_end + 1, ...])
 
 
 def test_fill():