Merge pull request #3807 from airqo-platform/save-satellite-predictions

add code to save model predictions to BigQuery
airqo-platform · Nov 7, 2024 · 03a4385 · 03a4385
2 parents ee61112 + 271a2a9
commit 03a4385
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 9 deletions.
diff --git a/src/spatial/configure.py b/src/spatial/configure.py
@@ -21,6 +21,9 @@ class Config:
         "GOOGLE_APPLICATION_CREDENTIALS_EMAIL"
     )
     PROJECT_BUCKET = os.getenv("PROJECT_BUCKET")
+    BIGQUERY_SATELLITE_MODEL_PREDICTIONS = os.getenv(
+        "BIGQUERY_SATELLITE_MODEL_PREDICTIONS"
+    )
 
 
 class ProductionConfig(Config):

diff --git a/src/spatial/views/derived_pm2_5.py b/src/spatial/views/derived_pm2_5.py
@@ -10,6 +10,7 @@
 
 logging.basicConfig(level=logging.ERROR)
 
+
 class PM25View:
     @staticmethod
     def get_pm25():
@@ -54,7 +55,11 @@ def get_pm25():
 
         except Exception as e:
             logging.error("An error occurred: %s", str(e))
-            return jsonify({"error": "An internal error has occurred!"}), 500, {"Content-Type": "application/json"}
+            return (
+                jsonify({"error": "An internal error has occurred!"}),
+                500,
+                {"Content-Type": "application/json"},
+            )
 
 
 #            return jsonify({'error': 'An internal error occurred'}), 500, {'Content-Type': 'application/json'}

diff --git a/src/spatial/views/satellite_predictions.py b/src/spatial/views/satellite_predictions.py
@@ -1,7 +1,9 @@
 from datetime import datetime, timezone
 
 import numpy as np
+import pandas as pd
 from flask import request, jsonify
+from google.oauth2 import service_account
 
 from configure import get_trained_model_from_gcs, Config, satellite_collections
 from models.SatellitePredictionModel import SatellitePredictionModel
@@ -38,15 +40,30 @@ def make_predictions():
             )
 
             prediction = model.predict(feature_array)[0]
-            return jsonify(
-                {
-                    "pm2_5_prediction": round(float(prediction), 3),
-                    "latitude": latitude,
-                    "longitude": longitude,
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                }
-            )
 
+            result = {
+                "pm2_5_prediction": round(float(prediction), 3),
+                "latitude": latitude,
+                "longitude": longitude,
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+            }
+            try:
+                df = pd.DataFrame([result])
+                credentials = service_account.Credentials.from_service_account_file(
+                    Config.CREDENTIALS
+                )
+
+                df.to_gbq(
+                    destination_table=f"{Config.BIGQUERY_SATELLITE_MODEL_PREDICTIONS}",
+                    project_id=Config.GOOGLE_CLOUD_PROJECT_ID,
+                    if_exists="append",
+                    credentials=credentials,
+                )
+
+            except Exception as e:
+                print(f"Error saving predictions to BigQuery: {e}")
+
+            return jsonify(result), 200
         except Exception as e:
             print(f"Error making predictions: {e}")
             return jsonify({"error": "An internal error has occurred"}), 500