Update app.py

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-import os
+from huggingface_hub import from_pretrained_keras
 
-# 1. Load the Model
-# We expect the .keras file to be in the same directory
-model = tf.keras.models.load_model("hk_transit_flow_net.keras")
+# 1. Download the Model from your Repository
+# This connects this Space to your uploaded model
+print("Downloading model...")
+model = from_pretrained_keras("WheelsTransit/HK-TransitFlow-Net")
 
 # Helper to map day names to integers
 DAY_MAP = {
@@ -15,13 +16,11 @@ DAY_MAP = {
 
 def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
     try:
-        # 1. Prepare Inputs
-        # We must match the exact shape and types used in training
-        
         # Handle empty route
         if not route_id or route_id.strip() == "":
             route_id = "UNKNOWN"
             
+        # Prepare inputs exactly as the model expects
         inputs = {
             'distance': np.array([[float(distance_meters)]]),
             'num_stops': np.array([[float(num_stops)]]),
@@ -30,11 +29,10 @@ def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
             'route_id': tf.constant([[str(route_id)]], dtype=tf.string)
         }
         
-        # 2. Run Prediction
+        # Run Prediction
         prediction = model.predict(inputs, verbose=0)
         seconds = float(prediction[0][0])
         
-        # 3. Format Output
         minutes = int(seconds // 60)
         rem_seconds = int(seconds % 60)
         
@@ -43,7 +41,7 @@ def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# 3. Build the Interface
+# Build the Interface
 iface = gr.Interface(
     fn=predict_eta,
     inputs=[
@@ -51,22 +49,12 @@ iface = gr.Interface(
         gr.Number(label="Number of Stops", value=10),
         gr.Slider(minimum=0, maximum=23, step=1, label="Hour of Day (0-23)", value=9),
         gr.Dropdown(choices=list(DAY_MAP.keys()), label="Day of Week", value="Monday"),
-        gr.Textbox(label="Route ID (Optional)", placeholder="e.g. 968+1+Yuen Long+Tin Hau", value="UNKNOWN")
+        gr.Textbox(label="Route ID", placeholder="968+1+...", value="UNKNOWN")
     ],
     outputs="text",
-    title="HK-TransitFlow-Net 🚌",
-    description="""
-    **Hong Kong Bus ETA Predictor**
-    
-    This model uses Deep Learning to predict bus travel time based on distance, stops, and time context.
-    
-    *   **Distance:** Physical distance of the path in meters.
-    *   **Route ID:** Internal ID (e.g., `968+1+...`). If unknown, leave as UNKNOWN.
-    *   **Note:** Trained on KMB & CTB data.
-    """,
+    title="HK-TransitFlow-Net Demo",
+    description="Live inference for HK Bus ETA prediction.",
     theme="soft"
 )
 
-# 4. Launch
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()