Fix missing utils package for model_tests

FR-DC · Oct 20, 2023 · 549da2b · 549da2b
1 parent 126417f
commit 549da2b
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 78 deletions.
diff --git a/pipeline/model_tests/chestnut_dec_may/evaluate.py b/pipeline/model_tests/chestnut_dec_may/evaluate.py
@@ -8,7 +8,7 @@
 from frdc.train import FRDCDataModule
 from frdc.train import FRDCModule
 from pipeline.model_tests.chestnut_dec_may.preprocess import preprocess
-from pipeline.model_tests.chestnut_dec_may.utils import get_dataset
+from pipeline.model_tests.utils import get_dataset
 
 # Get our Test
 # TODO: Ideally, we should have a separate dataset for testing.

diff --git a/pipeline/model_tests/chestnut_dec_may/main.py b/pipeline/model_tests/chestnut_dec_may/main.py
@@ -18,7 +18,7 @@
 from frdc.train import FRDCDataModule, FRDCModule
 from pipeline.model_tests.chestnut_dec_may.augmentation import augmentation
 from pipeline.model_tests.chestnut_dec_may.preprocess import preprocess
-from pipeline.model_tests.chestnut_dec_may.utils import get_dataset
+from pipeline.model_tests.utils import get_dataset
 
 
 def train_val_test_split(x: TensorDataset) -> list[Dataset, Dataset, Dataset]:

diff --git a/pipeline/model_tests/chestnut_dec_may/utils.py b/pipeline/model_tests/chestnut_dec_may/utils.py
diff --git a/pipeline/model_tests/utils.py b/pipeline/model_tests/utils.py
@@ -0,0 +1,12 @@
+from frdc.load import FRDCDataset
+from frdc.preprocess.extract_segments import extract_segments_from_bounds
+
+BANDS = ["NB", "NG", "NR", "RE", "NIR"]
+
+
+def get_dataset(site, date, version):
+    ds = FRDCDataset(site=site, date=date, version=version)
+    ar, order = ds.get_ar_bands(BANDS)
+    bounds, labels = ds.get_bounds_and_labels()
+    segments = extract_segments_from_bounds(ar, bounds, cropped=True)
+    return segments, labels