Merge pull request #86 from WorldCereal/move-prepdataframe

Isolated `prep_dataframe` method

Author: kvantricht

presto/eval.py

@@ -30,7 +30,7 @@
     get_sinusoid_encoding_table,
     param_groups_lrd,
 )
-from .utils import DEFAULT_SEED, device
+from .utils import DEFAULT_SEED, device, prep_dataframe
 
 logger = logging.getLogger("__main__")
 
@@ -73,9 +73,9 @@ def __init__(
         self.target_function = target_function
 
         train_data, val_data = WorldCerealLabelledDataset.split_df(train_data, val_size=val_size)
-        self.train_df = self.prep_dataframe(train_data, filter_function, dekadal=dekadal)
-        self.val_df = self.prep_dataframe(val_data, filter_function, dekadal=dekadal)
-        self.test_df = self.prep_dataframe(test_data, filter_function, dekadal=dekadal)
+        self.train_df = prep_dataframe(train_data, filter_function, dekadal=dekadal)
+        self.val_df = prep_dataframe(val_data, filter_function, dekadal=dekadal)
+        self.test_df = prep_dataframe(test_data, filter_function, dekadal=dekadal)
 
         self.spatial_inference_savedir = spatial_inference_savedir
 
@@ -90,23 +90,6 @@ def __init__(
         self.dekadal = dekadal
         self.ds_class = WorldCerealLabelled10DDataset if dekadal else WorldCerealLabelledDataset
 
-    @staticmethod
-    def prep_dataframe(
-        df: pd.DataFrame,
-        filter_function: Optional[Callable[[pd.DataFrame], pd.DataFrame]] = None,
-        dekadal: bool = False,
-    ):
-        # SAR cannot equal 0.0 since we take the log of it
-        cols = [f"SAR-{s}-ts{t}-20m" for s in ["VV", "VH"] for t in range(36 if dekadal else 12)]
-
-        df = df.drop_duplicates(subset=["sample_id", "lat", "lon", "end_date"])
-        df = df[~pd.isna(df).any(axis=1)]
-        df = df[~(df.loc[:, cols] == 0.0).any(axis=1)]
-        df = df.set_index("sample_id")
-        if filter_function is not None:
-            df = filter_function(df)
-        return df
-
     def _construct_finetuning_model(self, pretrained_model: Presto) -> PrestoFineTuningModel:
         model: PrestoFineTuningModel = cast(Callable, pretrained_model.construct_finetuning_model)(
             num_outputs=self.num_outputs

presto/inference.py

@@ -16,10 +16,9 @@
     DynamicWorld2020_2021,
 )
 from .dataset import WorldCerealBase
-from .eval import WorldCerealEval
 from .masking import BAND_EXPANSION
 from .presto import Presto
-from .utils import device
+from .utils import device, prep_dataframe
 
 # Index to band groups mapping
 IDX_TO_BAND_GROUPS = {
@@ -402,6 +401,6 @@ def process_parquet(df: pd.DataFrame) -> pd.DataFrame:
     df_pivot["end_date"] = df_pivot["end_date"].dt.date.astype(str)
     df_pivot["valid_date"] = df_pivot["valid_date"].dt.date.astype(str)
 
-    df_pivot = WorldCerealEval.prep_dataframe(df_pivot)
+    df_pivot = prep_dataframe(df_pivot)
 
     return df_pivot

presto/utils.py

@@ -235,7 +235,8 @@ def plot_for_group(grp_df):
                     partial(plot_map, diff_country, vmin=-1, cmap="coolwarm"),
                 ),
                 plot(
-                    f"{name} AEZ - CatBoost", partial(plot_map, diff_aez, vmin=-1, cmap="coolwarm")
+                    f"{name} AEZ - CatBoost",
+                    partial(plot_map, diff_aez, vmin=-1, cmap="coolwarm"),
                 ),
                 plot(
                     f"{name} Year - CatBoost",
@@ -308,3 +309,23 @@ def load_world_df() -> pd.DataFrame:
     world_df = gpd.read_file(data_dir / filename)
     world_df = world_df.drop(columns=["status", "color_code", "iso_3166_1_"])
     return world_df
+
+
+def prep_dataframe(
+    df: pd.DataFrame,
+    filter_function: Optional[Callable[[pd.DataFrame], pd.DataFrame]] = None,
+    dekadal: bool = False,
+):
+    """Duplication from eval.py but otherwise we would need catboost during
+    presto inference on OpenEO.
+    """
+    # SAR cannot equal 0.0 since we take the log of it
+    cols = [f"SAR-{s}-ts{t}-20m" for s in ["VV", "VH"] for t in range(36 if dekadal else 12)]
+
+    df = df.drop_duplicates(subset=["sample_id", "lat", "lon", "end_date"])
+    df = df[~pd.isna(df).any(axis=1)]
+    df = df[~(df.loc[:, cols] == 0.0).any(axis=1)]
+    df = df.set_index("sample_id")
+    if filter_function is not None:
+        df = filter_function(df)
+    return df

setup.py

@@ -32,7 +32,7 @@ def load_dependencies(tag: str) -> List[str]:
     long_description_content_type="text/markdown",
     author="Gabriel Tseng",
     author_email="[email protected]",
-    version="0.1.1",
+    version="0.1.2",
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: Other/Proprietary License",