Update app.py

app.py

@@ -4,23 +4,33 @@ import numpy as np
 import requests
 from huggingface_hub import hf_hub_download
 
-# --- 1. Fetch Route Options for Dropdown ---
-print("Fetching route list for dropdown...")
+# --- 1. Fetch Route Options (KMB/CTB Only) ---
+print("Fetching and filtering route list...")
 try:
     resp = requests.get("https://hkbus.github.io/hk-bus-crawling/routeFareList.min.json")
     if resp.status_code == 200:
         data = resp.json()
-        # Get all keys (e.g. "968+1+Yuen Long+Tin Hau")
-        # We sort them so they are easy to find
-        all_routes = sorted(list(data['routeList'].keys()))
+        raw_list = data['routeList']
+        
+        # Filter Logic: Only keep KMB or CTB
+        valid_companies = ['kmb', 'ctb']
+        filtered_routes = []
+        
+        for key, info in raw_list.items():
+            # Check if company exists and is in our allowed list
+            if 'co' in info and len(info['co']) > 0:
+                company = info['co'][0]
+                if company in valid_companies:
+                    filtered_routes.append(key)
+                    
+        all_routes = sorted(filtered_routes)
     else:
         all_routes = []
 except Exception as e:
     print(f"Error fetching routes: {e}")
     all_routes = []
 
-# Add "UNKNOWN" as the first/default option
-route_choices = ["UNKNOWN"] + all_routes
+print(f"Loaded {len(all_routes)} valid KMB/CTB routes.")
 
 # --- 2. Download and Load Model ---
 print("Downloading model...")
@@ -35,13 +45,29 @@ DAY_MAP = {
     "Thursday": 4, "Friday": 5, "Saturday": 6
 }
 
+# --- Helper: Search Logic ---
+def filter_routes(search_text):
+    """
+    Returns a list of routes matching the search text.
+    Limits to 100 results to prevent browser crash.
+    """
+    if not search_text:
+        return gr.Dropdown(choices=["UNKNOWN"] + all_routes[:20]) # Default top 20
+        
+    search_text = search_text.lower()
+    # Filter list
+    filtered = [r for r in all_routes if search_text in r.lower()]
+    
+    # Cap at 100 results
+    return gr.Dropdown(choices=["UNKNOWN"] + filtered[:100], value="UNKNOWN")
+
+
+# --- Prediction Logic ---
 def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
     try:
-        # 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)]]),
@@ -50,7 +76,6 @@ def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
             'route_id': tf.constant([[str(route_id)]], dtype=tf.string)
         }
         
-        # Run Prediction
         prediction = model.predict(inputs, verbose=0)
         seconds = float(prediction[0][0])
         
@@ -65,32 +90,32 @@ def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
 # --- 3. Build the UI ---
 with gr.Blocks() as demo:
     gr.Markdown("# HK-TransitFlow-Net Demo 🚌")
-    gr.Markdown("Live inference for HK Bus ETA prediction.")
+    gr.Markdown("Live inference for **KMB & CTB** Bus ETA prediction.")
     
     with gr.Row():
         with gr.Column():
-            # Inputs
+            gr.Markdown("### 1. Trip Details")
             dist_input = gr.Number(label="Distance (meters)", value=5000)
             stops_input = gr.Number(label="Number of Stops", value=10)
             hour_input = gr.Slider(minimum=0, maximum=23, step=1, label="Hour of Day (0-23)", value=9)
             day_input = gr.Dropdown(choices=list(DAY_MAP.keys()), label="Day of Week", value="Monday")
             
-            # UPDATED: Dropdown with search capabilities
-            route_input = gr.Dropdown(
-                choices=route_choices, 
-                label="Route ID", 
-                value="UNKNOWN",
-                filterable=True, # Allows typing to search
-                interactive=True
-            )
+        with gr.Column():
+            gr.Markdown("### 2. Route Selection")
+            gr.Markdown("*Type in the box below to find your route (e.g. '968')*")
             
-            predict_btn = gr.Button("Predict ETA", variant="primary")
+            # Search Box
+            route_search = gr.Textbox(label="Search Route Number", placeholder="Type route number...")
             
-        with gr.Column():
-            # Output
+            # Dropdown
+            route_input = gr.Dropdown(label="Select Route ID", choices=["UNKNOWN"], value="UNKNOWN", interactive=True)
+            
+            predict_btn = gr.Button("Predict ETA", variant="primary")
             output_text = gr.Textbox(label="Estimated Travel Time", lines=1)
 
-    # Event Listener
+    # --- Interaction Logic ---
+    route_search.change(fn=filter_routes, inputs=route_search, outputs=route_input)
+
     predict_btn.click(
         fn=predict_eta,
         inputs=[dist_input, stops_input, hour_input, day_input, route_input],