Update load.py

load.py

@@ -17,6 +17,39 @@ hft.download_model("sharktide", "HurricaneTrustNet")
 hft.download_model("sharktide", "TornadoNet")
 hft.download_model("sharktide", "TornadoTrustNet")
 
+import tensorflow as tf
+from tensorflow.keras import layers, models, callbacks
+from tensorflow.keras.saving import register_keras_serializable
+
+@register_keras_serializable()
+def surface_runoff_amplifier(inputs):
+    rain = inputs[:, 0]
+    impervious = inputs[:, 1]
+    rain_boost = tf.sigmoid((rain - 60) * 0.06)
+    impervious_boost = tf.sigmoid((impervious - 0.6) * 10)
+    return (1.0 + 0.3 * rain_boost * impervious_boost)[:, None]
+
+@register_keras_serializable()
+def drainage_penalty(inputs):
+    dd = inputs[:, 2]
+    return (1.0 - 0.4 * tf.sigmoid((dd - 3.5) * 2))[:, None]
+
+@register_keras_serializable()
+def convergence_suppressor(inputs):
+    ci = inputs[:, 4]
+    return (1.0 + 0.3 * tf.sigmoid((ci - 0.5) * 8))[:, None]
+
+@register_keras_serializable()
+def clip_modulation(x):
+    return tf.clip_by_value(x, 0.7, 1.3)
+
+CUSTOM_OBJECTS = {
+    'drainage_penalty': drainage_penalty,
+    'convergence_suppressor': convergence_suppressor,
+    'surface_runoff_amplifier': surface_runoff_amplifier,
+    'clip_modulation': clip_modulation
+}
+
 FireNet = hft.load_model("sharktide", "FireNet", "tf_model.h5", True)
 FireTrustNet = hft.load_model("sharktide", "FireTrustNet", "tf_model.h5", True)
 FireScaler = joblib.load("scalers/firetrust_scaler.pkl")
@@ -25,7 +58,8 @@ FloodNet = hft.load_model("sharktide", "FV-FloodNet", "tf_model.h5", True)
 FloodTrustNet = hft.load_model("sharktide", "FV-FloodTrustNet", "tf_model.h5", True)
 FloodScaler = joblib.load("scalers/FV-floodtrust_scaler.pkl")
 
-PV_FloodNet = hft.load_model("sharktide", "PV-FloodNet", "tf_model.h5", True)
+get_path = lambda usr, model: (str(hft.get_model_folder(user, model)) + "tf_model.h5")
+PV_FloodNet = tf.keras.models.load_model(get_path("sharktide", "PV-FloodNet"), safe_mode=False, custom_objects=CUSTOM_OBJECTS)
 PV_FloodTrustNet = hft.load_model("sharktide", "PV-FloodTrustNet", "tf_model.h5", True)
 PV_FloodScaler = joblib.load("scalers/PV-floodtrust_scaler.pkl")