Update benchmark.py

benchmark.py

@@ -3,15 +3,16 @@ import pathlib
 import opensr_test
 import matplotlib.pyplot as plt
 
-from typing import Callable
+from typing import Callable, Union
 
 
 def create_geotiff(
     model: Callable,
     fn: Callable,
-    datasets: list,
+    datasets: Union[str, list],
     output_path: str,
-    force: bool = False
+    force: bool = False,
+    **kwargs
 ) -> None:
     """Create all the GeoTIFFs for a specific dataset snippet 
 
@@ -26,13 +27,18 @@ def create_geotiff(
         force (bool, optional): If True, the dataset is redownloaded. Defaults 
             to False.
     """
+    
+    if datasets == "all":
+        datasets = opensr_test.datasets 
+
     for snippet in datasets:
         create_geotiff_batch(
             model=model,
             fn=fn,
             snippet=snippet,
             output_path=output_path,
-            force=force
+            force=force,
+            **kwargs
         )    
 
     return None
@@ -42,7 +48,8 @@ def create_geotiff_batch(
     fn: Callable,
     snippet: str,
     output_path: str,
-    force: bool = False
+    force: bool = False,
+    **kwargs
 ) -> pathlib.Path:
     """Create all the GeoTIFFs for a specific dataset snippet 
 
@@ -72,7 +79,7 @@ def create_geotiff_batch(
     output_path_dataset_png.mkdir(parents=True, exist_ok=True)
 
     # Load the dataset 
-    dataset = opensr_test.load(snippet, force=False)
+    dataset = opensr_test.load(snippet, force=force)
     lr_dataset, hr_dataset, metadata = dataset["L2A"], dataset["HRharm"], dataset["metadata"]
     for index in range(len(lr_dataset)):
         print(f"Processing {index}/{len(lr_dataset)}")
@@ -81,7 +88,8 @@ def create_geotiff_batch(
         results = fn(
             model=model,
             lr=lr_dataset[index],
-            hr=hr_dataset[index]
+            hr=hr_dataset[index],
+            **kwargs
         )
 
         # Get the image name
@@ -165,4 +173,6 @@ def plot(
         pathlib.Path: The output path where the plots and tables are 
         saved.
     """
-    pass
+    pass
+
+