|
import os |
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' |
|
import pickle |
|
|
|
import numpy as np |
|
import pandas as pd |
|
import plotly.express as px |
|
import streamlit as st |
|
import tensorflow as tf |
|
import tensorflow_hub as hub |
|
|
|
from sklearn.cluster import DBSCAN |
|
|
|
|
|
def read_stops(p: str): |
|
""" |
|
DOCSTRING |
|
""" |
|
return pd.read_csv(p) |
|
|
|
|
|
def read_encodings(p: str) -> tf.Tensor: |
|
""" |
|
Unpickle the Universal Sentence Encoder v4 encodings |
|
and return them |
|
|
|
This function doesn't make any attempt to patch the security holes in `pickle`. |
|
|
|
:param p: Path to the encodings |
|
|
|
:returns: A Tensor of the encodings with shape (number of sentences, 512) |
|
""" |
|
with open(p, 'rb') as f: |
|
encodings = pickle.load(f) |
|
return encodings |
|
|
|
|
|
def cluster_encodings(encodings: tf.Tensor) -> np.ndarray: |
|
""" |
|
DOCSTRING |
|
""" |
|
|
|
clusterer = DBSCAN(eps=0.7, min_samples=100).fit(encodings) |
|
return clusterer.labels_ |
|
|
|
|
|
def cluster_lat_lon(df: pd.DataFrame) -> np.ndarray: |
|
""" |
|
DOCSTRING |
|
""" |
|
|
|
clusterer = DBSCAN(eps=0.025, min_samples=100).fit(df[['latitude', 'longitude']]) |
|
return clusterer.labels_ |
|
|
|
|
|
def plot_example(df: pd.DataFrame, labels: np.ndarray) -> px.Figure: |
|
""" |
|
DOCSTRING |
|
""" |
|
plot_size = 800 |
|
labels = labels.astype('str') |
|
|
|
fig = px.scatter(df, x='longitude', y='latitude', |
|
hover_name='display_name', |
|
color=labels, |
|
opacity=0.5, |
|
color_discrete_sequence=px.colors.qualitative.Safe, |
|
template='presentation', |
|
width=plot_size, |
|
height=plot_size) |
|
|
|
return fig |
|
|
|
|
|
def plot_venice_blvd(df: pd.DataFrame, labels: np.ndarray) -> px.Figure: |
|
""" |
|
DOCSTRING |
|
""" |
|
px.set_mapbox_access_token(st.secrets['mapbox_token']) |
|
venice_blvd = {'lat': 34.008350, |
|
'lon': -118.425362} |
|
labels = labels.astype('str') |
|
|
|
fig = px.scatter_mapbox(df, lat='latitude', lon='longitude', |
|
color=labels, |
|
hover_name='display_name', |
|
center=venice_blvd, |
|
zoom=12, |
|
color_discrete_sequence=px.colors.qualitative.Dark24) |
|
|
|
|
|
return fig |
|
|
|
|
|
def main(data_path: str, enc_path: str): |
|
df = read_stops(data_path) |
|
|
|
|
|
example_labels = cluster_lat_lon(df) |
|
example_fig = plot_example(df, example_labels) |
|
|
|
|
|
encodings = read_encodings(enc_path) |
|
encoding_labels = cluster_encodings(encodings) |
|
venice_fig = plot_venice_blvd(df, encoding_labels) |
|
|
|
|
|
st.write('# Example of what DBSCAN does') |
|
st.plotly_chart(example_fig, use_container_width=True) |
|
|
|
st.write('# Venice Blvd') |
|
st.plotly_chart(example_fig, use_container_width=True) |
|
|
|
|
|
if __name__ == '__main__': |
|
import argparse |
|
|
|
p = argparse.ArgumentParser() |
|
p.add_argument('--data_path', |
|
nargs='?', |
|
default='data/stops.csv', |
|
help="Path to the dataset of LA Metro stops. Defaults to 'data/stops.csv'") |
|
p.add_argument('--enc_path', |
|
nargs='?', |
|
default='data/encodings.pkl', |
|
help="Path to the pickled encodings. Defaults to 'data/encodings.pkl'") |
|
args = p.parse_args() |
|
|
|
main(**vars(args)) |
|
|
|
|