Enable ee high-volume optional

awi-response · Nov 26, 2024 · fc58e08 · fc58e08
1 parent 9600e7f
commit fc58e08
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 8 deletions.
diff --git a/darts/src/darts/native.py b/darts/src/darts/native.py
@@ -46,6 +46,7 @@ def run_native_planet_pipeline(
     notcvis_model_name: str = "RTS_v6_notcvis.pt",
     device: Literal["cuda", "cpu", "auto"] | int | None = None,
     ee_project: str | None = None,
+    ee_use_highvolume: bool = True,
     patch_size: int = 1024,
     overlap: int = 16,
     batch_size: int = 8,
@@ -74,6 +75,7 @@ def run_native_planet_pipeline(
             Defaults to "cuda" if available, else "cpu".
         ee_project (str, optional): The Earth Engine project ID or number to use. May be omitted if
             project is defined within persistent API credentials obtained via `earthengine authenticate`.
+        ee_use_highvolume (bool, optional): Whether to use the high volume server (https://earthengine-highvolume.googleapis.com).
         patch_size (int, optional): The patch size to use for inference. Defaults to 1024.
         overlap (int, optional): The overlap to use for inference. Defaults to 16.
         batch_size (int, optional): The batch size to use for inference. Defaults to 8.
@@ -169,7 +171,7 @@ def run_native_planet_pipeline(
 
     debug_info()
     device = decide_device(device)
-    init_ee(ee_project)
+    init_ee(ee_project, ee_use_highvolume)
 
     # Find all PlanetScope orthotiles
     for fpath, outpath in planet_file_generator(orthotiles_dir, scenes_dir, output_data_dir):
@@ -219,6 +221,7 @@ def run_native_planet_pipeline_fast(
     notcvis_model_name: str = "RTS_v6_notcvis.pt",
     device: Literal["cuda", "cpu", "auto"] | int | None = None,
     ee_project: str | None = None,
+    ee_use_highvolume: bool = True,
     tpi_outer_radius: int = 100,
     tpi_inner_radius: int = 0,
     patch_size: int = 1024,
@@ -251,6 +254,7 @@ def run_native_planet_pipeline_fast(
             Defaults to "cuda" if available, else "cpu".
         ee_project (str, optional): The Earth Engine project ID or number to use. May be omitted if
             project is defined within persistent API credentials obtained via `earthengine authenticate`.
+        ee_use_highvolume (bool, optional): Whether to use the high volume server (https://earthengine-highvolume.googleapis.com).
         tpi_outer_radius (int, optional): The outer radius of the annulus kernel for the tpi calculation
             in m. Defaults 100m.
         tpi_inner_radius (int, optional): The inner radius of the annulus kernel for the tpi calculation
@@ -286,7 +290,7 @@ def run_native_planet_pipeline_fast(
 
     debug_info()
     device = decide_device(device)
-    init_ee(ee_project)
+    init_ee(ee_project, ee_use_highvolume)
 
     client = Client()
     logger.info(f"Using Dask client: {client}")
@@ -351,6 +355,7 @@ def run_native_sentinel2_pipeline(
     notcvis_model_name: str = "RTS_v6_notcvis_s2native.pt",
     device: Literal["cuda", "cpu", "auto"] | int | None = None,
     ee_project: str | None = None,
+    ee_use_highvolume: bool = True,
     patch_size: int = 1024,
     overlap: int = 16,
     batch_size: int = 8,
@@ -378,6 +383,7 @@ def run_native_sentinel2_pipeline(
             Defaults to "cuda" if available, else "cpu".
         ee_project (str, optional): The Earth Engine project ID or number to use. May be omitted if
             project is defined within persistent API credentials obtained via `earthengine authenticate`.
+        ee_use_highvolume (bool, optional): Whether to use the high volume server (https://earthengine-highvolume.googleapis.com).
         patch_size (int, optional): The patch size to use for inference. Defaults to 1024.
         overlap (int, optional): The overlap to use for inference. Defaults to 16.
         batch_size (int, optional): The batch size to use for inference. Defaults to 8.
@@ -435,7 +441,7 @@ def run_native_sentinel2_pipeline(
 
     debug_info()
     device = decide_device(device)
-    init_ee(ee_project)
+    init_ee(ee_project, ee_use_highvolume)
 
     # Find all Sentinel 2 scenes
     for fpath in sentinel2_dir.glob("*/"):
@@ -487,6 +493,7 @@ def run_native_sentinel2_pipeline_fast(
     notcvis_model_name: str = "RTS_v6_notcvis_s2native.pt",
     device: Literal["cuda", "cpu", "auto"] | int | None = None,
     ee_project: str | None = None,
+    ee_use_highvolume: bool = True,
     tpi_outer_radius: int = 100,
     tpi_inner_radius: int = 0,
     patch_size: int = 1024,
@@ -519,6 +526,7 @@ def run_native_sentinel2_pipeline_fast(
             Defaults to "cuda" if available, else "cpu".
         ee_project (str, optional): The Earth Engine project ID or number to use. May be omitted if
             project is defined within persistent API credentials obtained via `earthengine authenticate`.
+        ee_use_highvolume (bool, optional): Whether to use the high volume server (https://earthengine-highvolume.googleapis.com).
         tpi_outer_radius (int, optional): The outer radius of the annulus kernel for the tpi calculation
             in m. Defaults to 100m.
         tpi_inner_radius (int, optional): The inner radius of the annulus kernel for the tpi calculation
@@ -554,7 +562,7 @@ def run_native_sentinel2_pipeline_fast(
 
     debug_info()
     device = decide_device(device)
-    init_ee(ee_project)
+    init_ee(ee_project, ee_use_highvolume)
 
     client = Client()
     logger.info(f"Using Dask client: {client}")

diff --git a/darts/src/darts/utils/earthengine.py b/darts/src/darts/utils/earthengine.py
@@ -10,20 +10,22 @@
 logger = logging.getLogger(__name__)
 
 
-def init_ee(project: str | None = None) -> None:
+def init_ee(project: str | None = None, use_highvolume: bool = True) -> None:
     """Initialize Earth Engine. Authenticate if necessary.
 
     Args:
         project (str): The project name.
+        use_highvolume (bool): Whether to use the high volume server (https://earthengine-highvolume.googleapis.com).
 
     """
-    logger.debug(f"Initializing Earth Engine with project {project}")
+    logger.debug(f"Initializing Earth Engine with project {project} {'with high volume' if use_highvolume else ''}")
+    opt_url = "https://earthengine-highvolume.googleapis.com" if use_highvolume else None
     try:
-        ee.Initialize(project=project)
+        ee.Initialize(project=project, opt_url=opt_url)
         # geemap.ee_initialize(project=project, opt_url="https://earthengine-highvolume.googleapis.com")
     except Exception:
         logger.debug("Initializing Earth Engine failed, trying to authenticate before")
         ee.Authenticate()
-        ee.Initialize(project=project)
+        ee.Initialize(project=project, opt_url=opt_url)
         # geemap.ee_initialize(project=project, opt_url="https://earthengine-highvolume.googleapis.com")
     logger.debug("Earth Engine initialized")