From 7229437a5e676930b7e0031d4c03fc1df0606885 Mon Sep 17 00:00:00 2001
From: Joseph Xu <jzxu@google.com>
Date: Thu, 16 Jan 2025 14:24:17 -0800
Subject: [PATCH] Add cloud project to EE initialization.

PiperOrigin-RevId: 716378472
---
 src/skai/earth_engine.py      | 9 ++++++---
 src/skai/generate_examples.py | 4 +++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/skai/earth_engine.py b/src/skai/earth_engine.py
index f12d3232..4d2a57f3 100644
--- a/src/skai/earth_engine.py
+++ b/src/skai/earth_engine.py
@@ -156,7 +156,9 @@ def get_open_buildings(regions: list[ShapelyGeometry],
     )
 
 
-def initialize(service_account: str, private_key: str | None) -> bool:
+def initialize(
+    cloud_project: str, service_account: str, private_key: str | None
+) -> bool:
   """Initializes EE server connection.
 
   When not using a service account, this function assumes that the user has
@@ -164,6 +166,7 @@ def initialize(service_account: str, private_key: str | None) -> bool:
   "earthengine authenticate".
 
   Args:
+    cloud_project: Google Cloud project id.
     service_account: If not empty, the service account to use. Otherwise
       defaults to caller's personal account.
     private_key: Private key for service account.
@@ -174,9 +177,9 @@ def initialize(service_account: str, private_key: str | None) -> bool:
   try:
     if service_account:
       credentials = ee.ServiceAccountCredentials(service_account, private_key)
-      ee.Initialize(credentials)
+      ee.Initialize(credentials=credentials, project=cloud_project)
     else:
-      ee.Initialize()
+      ee.Initialize(project=cloud_project)
   except ee.EEException as e:
     logging.error('Error initializing Earth Engine: %s', e)
     return False
diff --git a/src/skai/generate_examples.py b/src/skai/generate_examples.py
index 7c6c30c9..9f3986db 100644
--- a/src/skai/generate_examples.py
+++ b/src/skai/generate_examples.py
@@ -330,7 +330,9 @@ def download_building_footprints(
     )
   elif config.buildings_method == 'open_buildings':
     if not earth_engine.initialize(
-        config.earth_engine_service_account, config.earth_engine_private_key
+        config.cloud_project,
+        config.earth_engine_service_account,
+        config.earth_engine_private_key,
     ):
       raise NotInitializedEarthEngineError()
     logging.info('Querying Open Buildings centroids. This may take a while.')