gtfs-viewer-app.py

# # GTFS Viewer

# ## Refs
# - https://gtfs.org/ja/
# - https://www.gtfs.jp/
# - https://nttdocomo-developers.jp/entry/20231218_1

# ## Setup
# $ pip install streamlit streamlit-folium folium pydeck pandas requests geopy gtfs-realtime-bindings
# $ streamlit run gtfs-viewer-app.py

# Or get an access token from here https://developer.odpt.org/
# $ ACCESS_TOKEN=<your_access_token> streamlit run gtfs-viewer-app.py

import io
import os
import time
import urllib.error
import urllib.request
import zipfile

import folium
import pandas as pd
import pydeck as pdk
import requests
import streamlit as st
from folium.features import Marker, PolyLine
from folium.plugins import BeautifyIcon
from geopy.distance import geodesic
from google.transit import gtfs_realtime_pb2
from streamlit_folium import st_folium

ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN")
DATA_DIR = os.environ.get("DATA_DIR", "./data")
GTFS_JP_DATA_DIR = os.path.join(DATA_DIR, "gtfs_jp")

CENTER = [35.681093, 139.767043]  # TOKYO STATION


def download_gtfs_jp():
    # See https://ckan.odpt.org/dataset/b_bus_gtfs_jp-toei for details
    print("download gtfs-jp data")
    try:
        os.makedirs(GTFS_JP_DATA_DIR)
    except FileExistsError:
        pass

    url = (
        f"https://api.odpt.org/api/v4/files/Toei/data/ToeiBus-GTFS.zip?acl:consumerKey={ACCESS_TOKEN}"
        if ACCESS_TOKEN
        else "https://api-public.odpt.org/api/v4/files/Toei/data/ToeiBus-GTFS.zip"
    )
    with (
        requests.get(url) as res,
        io.BytesIO(res.content) as bytes_io,
        zipfile.ZipFile(bytes_io) as zip,
    ):
        zip.extractall(GTFS_JP_DATA_DIR)


@st.cache_data
def load_stops_df():
    return pd.read_csv(
        os.path.join(GTFS_JP_DATA_DIR, "stops.txt"),
        usecols=["stop_id", "stop_name", "stop_lat", "stop_lon"],
    )


def filter_stops_df(stops_df, radius):
    return stops_df[
        stops_df.apply(
            lambda x: geodesic((x["stop_lat"], x["stop_lon"]), CENTER).m < radius,
            axis=1,
        )
    ]


@st.cache_resource(ttl=3600)
def download_gtfs_realtime(time_steps):
    # See https://ckan.odpt.org/dataset/b_bus_gtfs_rt-toei for details
    print(f"download gtf realtime data at time steps {time_steps}")
    url = (
        f"https://api.odpt.org/api/v4/gtfs/realtime/ToeiBus?acl:consumerKey={ACCESS_TOKEN}"
        if ACCESS_TOKEN
        else "https://api-public.odpt.org/api/v4/gtfs/realtime/ToeiBus"
    )
    feed = gtfs_realtime_pb2.FeedMessage()

    records = []
    with urllib.request.urlopen(url) as res:
        feed.ParseFromString(res.read())
        for entity in feed.entity:
            # print(entity)
            record = [
                entity.id,
                entity.vehicle.trip.trip_id,
                entity.vehicle.trip.route_id,
                entity.vehicle.trip.direction_id,
                entity.vehicle.position.latitude,
                entity.vehicle.position.longitude,
                entity.vehicle.current_stop_sequence,
                entity.vehicle.timestamp,
                entity.vehicle.stop_id,
            ]
            records.append(record)

    return pd.DataFrame(
        records,
        columns=[
            "id",
            "trip_id",
            "route_id",
            "direction_id",
            "lat",
            "lon",
            "current_stop_sequence",
            "timestamp",
            "stop_id",
        ],
    )


def filter_realtime_df(realtime_df, radius):
    return realtime_df[
        realtime_df.apply(
            lambda x: geodesic((x["lat"], x["lon"]), CENTER).m < radius,
            axis=1,
        )
    ]


def hex_to_rgb(h):
    h = h.lstrip("#")
    return tuple(int(h[i : i + 2], 16) for i in (0, 2, 4))


def create_folium_map(
    stops_df, realtime_df, previous_realtime_df, previous_previous_realtime_df
):
    m = folium.Map(location=CENTER, zoom_start=11)
    fg = folium.FeatureGroup(name="GTFS")

    if stops_df is not None:
        for index, row in stops_df.iterrows():
            marker = Marker(
                location=[row.stop_lat, row.stop_lon],
                popup=row.stop_name,
                icon=BeautifyIcon(
                    icon="tent",
                    border_color="#00ABDC",
                    background_color="#00ABDC",
                    iconSize=[10, 10],
                    iconAnchor=[5, 5],
                ),
            )
            fg.add_child(marker)

    previous_location_by_id = {}
    if previous_realtime_df is not None:
        for index, row in previous_realtime_df.iterrows():
            previous_location_by_id[row.id] = [row.lat, row.lon]

    previous_previous_location_by_id = {}
    if previous_previous_realtime_df is not None:
        for index, row in previous_previous_realtime_df.iterrows():
            previous_previous_location_by_id[row.id] = [row.lat, row.lon]

    if realtime_df is not None:
        for index, row in realtime_df.iterrows():
            current_location = [row.lat, row.lon]
            marker = Marker(
                location=current_location,
                popup=row.id,
                icon=BeautifyIcon(
                    icon="tent",
                    border_color="#0000FF",
                    background_color="#0000FF",
                    iconSize=[10, 10],
                    iconAnchor=[5, 5],
                ),
            )
            fg.add_child(marker)

            previous_location = previous_location_by_id.get(row.id)
            if previous_location:
                line = PolyLine(
                    [previous_location, current_location],
                    color="#0000FF",
                    weight=2.5,
                    opacity=1,
                )
                fg.add_child(line)

                previous_previous_location = previous_previous_location_by_id.get(
                    row.id
                )

                # print(row.id, previous_previous_location, previous_location, current_location)
                if previous_previous_location:
                    line = PolyLine(
                        [previous_previous_location, previous_location],
                        color="#8888FF",
                        weight=2.5,
                        opacity=1,
                    )
                    fg.add_child(line)

    return m, fg


def create_pydeck_map(stops_df, realtime_df):
    return pdk.Deck(
        map_style="mapbox://styles/mapbox/light-v9",
        initial_view_state=pdk.ViewState(
            latitude=CENTER[0],
            longitude=CENTER[1],
            zoom=11,
            pitch=0,
        ),
        layers=[
            pdk.Layer(
                "ScatterplotLayer",
                data=stops_df,
                get_position="[stop_lon, stop_lat]",
                get_color=hex_to_rgb("#00ABDC"),
                get_radius=50,
            ),
            pdk.Layer(
                "ScatterplotLayer",
                data=realtime_df,
                get_position="[lon, lat]",
                get_color=hex_to_rgb("#0000FF"),
                get_radius=50,
            ),
        ],
    )


def main():
    st.write("# GTFS Viewer")

    time_steps = st.session_state.get("time_steps", 0)

    gtfs_jp_data_downloaded = os.path.exists(GTFS_JP_DATA_DIR)

    with st.sidebar:

        def on_click():
            with st.spinner("Downloading..."):
                download_gtfs_jp()

        st.button(
            "Download GTFS-JP data",
            on_click=on_click,
            disabled=gtfs_jp_data_downloaded,
        )
        show_stops = st.toggle("Show stops", disabled=not gtfs_jp_data_downloaded)

        st.divider()

        show_realtime = st.toggle("Show realtime")

        refresh_interval = st.selectbox(
            "refresh interval",
            [30, 60, 300],
            index=0,
            disabled=not show_realtime,
        )

        st.divider()

        radius = st.selectbox(
            "filter radius",
            [1000, 5000, 100000],
            index=1,
        )

        map_type = st.selectbox(
            "map type",
            ["folium", "pydeck"],
            index=0,
        )

    if show_stops:
        stops_df = load_stops_df()
        filtered_stops_df = filter_stops_df(stops_df, radius)
    else:
        stops_df = filtered_stops_df = None

    if show_realtime:
        realtime_df = download_gtfs_realtime(time_steps)
        filtered_realtime_df = filter_realtime_df(realtime_df, radius)

        previous_time_steps = time_steps - refresh_interval
        previous_realtime_df = (
            download_gtfs_realtime(previous_time_steps)
            if previous_time_steps >= 0
            else None
        )

        previous_previous_time_steps = previous_time_steps - refresh_interval
        previous_previous_realtime_df = (
            download_gtfs_realtime(previous_previous_time_steps)
            if previous_previous_time_steps >= 0
            else None
        )

    else:
        realtime_df = filtered_realtime_df = None
        previous_realtime_df = None
        previous_previous_realtime_df = None

    if map_type == "folium":
        m, fg = create_folium_map(
            filtered_stops_df,
            filtered_realtime_df,
            previous_realtime_df,
            previous_previous_realtime_df,
        )
        st_folium(m, feature_group_to_add=fg, height=550)
    else:
        m = create_pydeck_map(
            filtered_stops_df, filtered_realtime_df
        )  # TODO: for previous_realtime_df
        st.pydeck_chart(m)

    if show_realtime:
        time.sleep(refresh_interval)
        st.session_state["time_steps"] = time_steps + refresh_interval
        st.rerun()


if __name__ == "__main__":
    main()