Update app.py

app.py

@@ -2,75 +2,191 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 import requests
+import math
 from huggingface_hub import hf_hub_download
 
-# --- 1. Fetch Route Options (KMB/CTB Only) ---
-print("Fetching and filtering route list...")
+# --- Global Data Storage ---
+ROUTE_DATA = {}
+STOP_DATA = {}
+ALL_ROUTE_KEYS = []
+
+# --- 1. Fetch Data (Routes & Stops) ---
+print("Fetching route and stop data...")
 try:
     resp = requests.get("https://hkbus.github.io/hk-bus-crawling/routeFareList.min.json")
     if resp.status_code == 200:
-        data = resp.json()
-        raw_list = data['routeList']
+        json_db = resp.json()
+        raw_routes = json_db['routeList']
+        STOP_DATA = json_db['stopList']
         
         # 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
+        for key, info in raw_routes.items():
             if 'co' in info and len(info['co']) > 0:
                 company = info['co'][0]
                 if company in valid_companies:
-                    filtered_routes.append(key)
+                    # Store the whole object so we can look up stops later
+                    ROUTE_DATA[key] = info
+                    ALL_ROUTE_KEYS.append(key)
                     
-        all_routes = sorted(filtered_routes)
+        ALL_ROUTE_KEYS.sort()
     else:
-        all_routes = []
+        print("Failed to download route data")
 except Exception as e:
-    print(f"Error fetching routes: {e}")
-    all_routes = []
+    print(f"Error fetching data: {e}")
 
-print(f"Loaded {len(all_routes)} valid KMB/CTB routes.")
+print(f"Loaded {len(ALL_ROUTE_KEYS)} valid KMB/CTB routes.")
 
 # --- 2. Download and Load Model ---
 print("Downloading model...")
-model_path = hf_hub_download(repo_id="WheelsTransit/HK-TransitFlow-Net", filename="hk_transit_flow_net.keras")
-
-print("Loading Keras model...")
-model = tf.keras.models.load_model(model_path)
+try:
+    model_path = hf_hub_download(repo_id="WheelsTransit/HK-TransitFlow-Net", filename="hk_transit_flow_net.keras")
+    print("Loading Keras model...")
+    model = tf.keras.models.load_model(model_path)
+except Exception as e:
+    print(f"Model load failed: {e}")
+    model = None
 
-# Helper to map day names to integers
+# --- Helpers ---
 DAY_MAP = {
     "Sunday": 0, "Monday": 1, "Tuesday": 2, "Wednesday": 3,
     "Thursday": 4, "Friday": 5, "Saturday": 6
 }
 
-# --- Helper: Search Logic ---
+def haversine_distance(coords):
+    """Calculates length of a coordinate list in meters."""
+    R = 6371000 
+    total_dist = 0
+    for i in range(len(coords) - 1):
+        lon1, lat1 = coords[i]
+        lon2, lat2 = coords[i+1]
+        dlon = math.radians(lon2 - lon1)
+        dlat = math.radians(lat2 - lat1)
+        a = math.sin(dlat/2)**2 + math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) * math.sin(dlon/2)**2
+        c = 2 * math.asin(math.sqrt(a))
+        total_dist += R * c
+    return total_dist
+
+# --- Dynamic UI Logic ---
+
 def filter_routes(search_text):
-    """
-    Returns a list of routes matching the search text.
-    Limits to 100 results to prevent browser crash.
-    """
+    """Filters the route dropdown based on search text."""
     if not search_text:
-        return gr.Dropdown(choices=["UNKNOWN"] + all_routes[:20]) # Default top 20
-        
+        return gr.Dropdown(choices=["UNKNOWN"] + ALL_ROUTE_KEYS[: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
+    filtered = [r for r in ALL_ROUTE_KEYS if search_text in r.lower()]
     return gr.Dropdown(choices=["UNKNOWN"] + filtered[:100], value="UNKNOWN")
 
+def update_stop_dropdowns(route_key):
+    """
+    When a route is selected, fetch its stops and populate the Start/End dropdowns.
+    Returns: (Start_Dropdown_Update, End_Dropdown_Update)
+    """
+    if not route_key or route_key == "UNKNOWN" or route_key not in ROUTE_DATA:
+        return gr.Dropdown(choices=[], value=None), gr.Dropdown(choices=[], value=None)
+    
+    route_info = ROUTE_DATA[route_key]
+    company = route_info['co'][0]
+    stop_ids = route_info['stops'].get(company, [])
+    
+    # Create readable names: "1. StopName (ID)"
+    stop_options = []
+    for idx, sid in enumerate(stop_ids):
+        # Fetch name from STOP_DATA
+        name_en = "Unknown"
+        if sid in STOP_DATA:
+            name_en = STOP_DATA[sid]['name']['en']
+        
+        label = f"{idx+1}. {name_en} ({sid})"
+        stop_options.append(label)
+        
+    return gr.Dropdown(choices=stop_options, value=None), gr.Dropdown(choices=stop_options, value=None)
+
+def calculate_real_metrics(route_key, start_str, end_str):
+    """
+    Downloads Waypoints and calculates actual distance/stops between two selected stops.
+    """
+    if route_key == "UNKNOWN" or not start_str or not end_str:
+        return None, None, "Please select a Route, Start Stop, and End Stop."
 
-# --- Prediction Logic ---
-def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
     try:
-        if not route_id or route_id.strip() == "":
-            route_id = "UNKNOWN"
+        # Extract Index from string "1. Name (ID)"
+        start_idx = int(start_str.split(".")[0]) - 1
+        end_idx = int(end_str.split(".")[0]) - 1
+        
+        if start_idx >= end_idx:
+            return None, None, "Error: Start Stop must be before End Stop."
+
+        # Fetch Route Info for GTFS ID
+        route_info = ROUTE_DATA[route_key]
+        gtfs_id = route_info.get('gtfsId')
+        company = route_info['co'][0]
+        bound = route_info['bound'].get(company)
+        
+        if not gtfs_id or not bound:
+            return None, None, "Error: No GTFS data for this route."
+
+        # Download Waypoints
+        url = f"https://hkbus.github.io/route-waypoints/{gtfs_id}-{bound}.json"
+        resp = requests.get(url)
+        if resp.status_code != 200:
+            return None, None, "Error: Could not download route path data."
             
+        geojson = resp.json()
+        features = geojson.get('features', [])
+        
+        # Determine Segments
+        # Logic: Feature[i] is path from Stop[i] to Stop[i+1]
+        segments = []
+        if features and features[0]['geometry']['type'] == 'MultiLineString':
+             segments = features[0]['geometry']['coordinates']
+        elif features:
+             segments = [f['geometry']['coordinates'] for f in features]
+             
+        # Sum distance for the specific range
+        total_dist = 0
+        # We need to sum segments from start_idx to end_idx - 1
+        # E.g. Stop 0 to Stop 2 requires Segment 0 (0->1) and Segment 1 (1->2)
+        
+        for i in range(start_idx, end_idx):
+            if i < len(segments):
+                total_dist += haversine_distance(segments[i])
+                
+        num_stops = end_idx - start_idx
+        
+        return total_dist, num_stops, None # None = No error
+
+    except Exception as e:
+        return None, None, f"Calculation Error: {str(e)}"
+
+# --- Prediction Logic ---
+
+def smart_predict(manual_dist, manual_stops, hour, day_name, route_id, start_stop, end_stop):
+    status_msg = ""
+    
+    # 1. Automatic Calculation Logic
+    if route_id != "UNKNOWN" and start_stop and end_stop:
+        calc_dist, calc_stops, error = calculate_real_metrics(route_id, start_stop, end_stop)
+        
+        if error:
+            status_msg = f"⚠️ {error} Using manual inputs."
+            final_dist = manual_dist
+            final_stops = manual_stops
+        else:
+            final_dist = calc_dist
+            final_stops = calc_stops
+            status_msg = f"✅ Calculated from map: {final_dist:.0f}m / {final_stops} stops."
+    else:
+        final_dist = manual_dist
+        final_stops = manual_stops
+        status_msg = "ℹ️ Using manual inputs."
+
+    # 2. Model Inference
+    try:
         inputs = {
-            'distance': np.array([[float(distance_meters)]]),
-            'num_stops': np.array([[float(num_stops)]]),
+            'distance': np.array([[float(final_dist)]]),
+            'num_stops': np.array([[float(final_stops)]]),
             'hour': np.array([[int(hour)]]),
             'day_of_week': np.array([[int(DAY_MAP[day_name])]]),
             'route_id': tf.constant([[str(route_id)]], dtype=tf.string)
@@ -82,46 +198,67 @@ def predict_eta(distance_meters, num_stops, hour, day_name, route_id):
         minutes = int(seconds // 60)
         rem_seconds = int(seconds % 60)
         
-        return f"{minutes} min {rem_seconds} sec ({seconds:.1f}s)"
+        result_str = f"⏱️ ETA: {minutes} min {rem_seconds} sec"
+        
+        # Return updated boxes + result
+        return final_dist, final_stops, f"{status_msg}\n\n{result_str}"
     
     except Exception as e:
-        return f"Error: {str(e)}"
+        return final_dist, final_stops, f"Model Error: {str(e)}"
+
 
 # --- 3. Build the UI ---
-with gr.Blocks() as demo:
+with gr.Blocks(title="HK-TransitFlow-Net") as demo:
     gr.Markdown("# HK-TransitFlow-Net Demo")
-    gr.Markdown("Live inference for **KMB & CTB** Bus ETA prediction.")
-    gr.Markdown("Predicts estimated journey time for a distance.")
-    gr.Markdown("Model URL: https://huggingface.co/WheelsTransit/HK-TransitFlow-Net")
+    gr.Markdown("Predicts KMB/CTB bus travel time. Select a route and stops to auto-calculate distance.")
     
     with gr.Row():
+        # Left Column: Inputs
         with gr.Column():
-            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")
+            gr.Markdown("### 1. Route Selection")
+            route_search = gr.Textbox(label="Search Route", placeholder="Type e.g. '968'")
+            route_dropdown = gr.Dropdown(label="Select Route ID", choices=["UNKNOWN"], value="UNKNOWN", interactive=True)
             
-        with gr.Column():
-            gr.Markdown("### 2. Route Selection (Optional - Select UNKNOWN to predict without route)")
-            gr.Markdown("*Type in the box below to find your route (e.g. '968')*")
-            
-            # Search Box
-            route_search = gr.Textbox(label="Search Route Number", placeholder="Type route number...")
+            with gr.Row():
+                start_dropdown = gr.Dropdown(label="Start Stop", choices=[], interactive=True)
+                end_dropdown = gr.Dropdown(label="End Stop", choices=[], interactive=True)
             
-            # Dropdown
-            route_input = gr.Dropdown(label="Select Route ID", choices=["UNKNOWN"], value="UNKNOWN", interactive=True)
+            gr.Markdown("---")
+            gr.Markdown("### 2. Time & Details")
+            with gr.Row():
+                hour_input = gr.Slider(minimum=0, maximum=23, step=1, label="Hour (0-23)", value=9)
+                day_input = gr.Dropdown(choices=list(DAY_MAP.keys()), label="Day", value="Monday")
             
+            with gr.Row():
+                # These update automatically if stops are picked
+                dist_input = gr.Number(label="Distance (m)", value=5000)
+                stops_input = gr.Number(label="Stops Count", value=10)
+
             predict_btn = gr.Button("Predict ETA", variant="primary")
-            output_text = gr.Textbox(label="Estimated Travel Time", lines=1)
 
-    # --- Interaction Logic ---
-    route_search.change(fn=filter_routes, inputs=route_search, outputs=route_input)
+        # Right Column: Output
+        with gr.Column():
+            gr.Markdown("### Result")
+            output_text = gr.Textbox(label="Prediction", lines=4)
+            gr.Markdown("*Note: If you select stops, distance is calculated automatically from the map.*")
 
+    # --- Event Wiring ---
+    
+    # 1. Search filter
+    route_search.change(fn=filter_routes, inputs=route_search, outputs=route_dropdown)
+    
+    # 2. Populate Stops when Route Selected
+    route_dropdown.change(
+        fn=update_stop_dropdowns, 
+        inputs=route_dropdown, 
+        outputs=[start_dropdown, end_dropdown]
+    )
+    
+    # 3. Predict Button
     predict_btn.click(
-        fn=predict_eta,
-        inputs=[dist_input, stops_input, hour_input, day_input, route_input],
-        outputs=output_text
+        fn=smart_predict,
+        inputs=[dist_input, stops_input, hour_input, day_input, route_dropdown, start_dropdown, end_dropdown],
+        outputs=[dist_input, stops_input, output_text] # Updates the number boxes too!
     )
 
 if __name__ == "__main__":