call_models/obs_map.py

from typing import Tuple
import logging

import pandas as pd
from datasets import load_dataset
import streamlit as st
import folium
from streamlit_folium import st_folium

import whale_viewer as sw_wv
from fix_tabrender import js_show_zeroheight_iframe

m_logger = logging.getLogger(__name__)
# we can set the log level locally for funcs in this module
#g_m_logger.setLevel(logging.DEBUG)
m_logger.setLevel(logging.INFO)

# TODO: refactor so we have richer data: a tuple or dict combining
# the dropdown label, the tileset name, the attribution - everything
# needed to make the map logic simplified
tile_sets = [
    'Open Street Map',
    #'Stamen Terrain',
    #'Stamen Toner',
    'Esri Ocean',
    'Esri Images',
    'Stamen Watercolor',
    'CartoDB Positron',
    #'CartoDB Dark_Matter'
]

# a list of unique colours for each whale class (for the map)
_colors = [
    "#FFD700", # Gold
    "#FF5733", # Red
    "#33FF57", # Green
    "#3357FF", # Blue
    "#FFFF33", # Yellow
    "#FF33FF", # Magenta
    "#33FFFF", # Cyan
    "#FF8C00", # Dark Orange
    "#8A2BE2", # Blue Violet
    "#DEB887", # Burlywood
    "#5F9EA0", # Cadet Blue
    "#D2691E", # Chocolate
    "#FF4500", # Orange Red
    "#2E8B57", # Sea Green
    "#DA70D6", # Orchid
    "#FF6347", # Tomato
    "#7FFF00", # Chartreuse
    "#DDA0DD", # Plum
    "#A0522D", # Sienna
    "#4682B4", # Steel Blue
    "#7B68EE", # Medium Slate Blue
    "#F0E68C", # Khaki
    "#B22222", # Firebrick
    "#FF1493", # Deep Pink
    "#FFFACD", # Lemon Chiffon
    "#20B2AA", # Light Sea Green
    "#778899" # Light Slate Gray
]

whale2color = {k: v for k, v in zip(sw_wv.WHALE_CLASSES, _colors)}

def create_map(tile_name:str, location:Tuple, zoom_start: int = 7):
    # https://xyzservices.readthedocs.io/en/stable/gallery.html 
    # get teh attribtuions from here once we pick the 2-3-4 options 
    # make esri ocean the default
    m = folium.Map(location=location, zoom_start=zoom_start,
                   tiles='Esri.OceanBasemap', attr="Esri")
    #m = folium.Map(location=location, zoom_start=zoom_start)

    attr = ""
    if tile_name == 'Open Street Map':
        folium.TileLayer('openstreetmap').add_to(m)
        pass

    #Esri.OceanBasemap
    elif tile_name == 'Esri Ocean':
        pass # made this one default ()
        #attr = "Esri"
        #folium.TileLayer('Esri.OceanBasemap', attr=attr).add_to(m)

    elif tile_name == 'Esri Images':
        attr = "Esri — Source: Esri, i-cubed, USDA"
        #folium.TileLayer('stamenterrain', attr=attr).add_to(m)
        folium.TileLayer('Esri.WorldImagery', attr=attr).add_to(m)
    elif tile_name == 'Stamen Toner':
        attr = "Stamen"
        folium.TileLayer('stamentoner', attr=attr).add_to(m)
    elif tile_name == 'Stamen Watercolor':
        attr = "Stamen"
        folium.TileLayer('Stadia.StamenWatercolor', attr=attr).add_to(m)
    elif tile_name == 'CartoDB Positron':
        folium.TileLayer('cartodb positron').add_to(m)
    elif tile_name == 'CartoDB Dark_Matter':
        folium.TileLayer('cartodb dark_matter').add_to(m)

    #folium.LayerControl().add_to(m)
    return m



def present_obs_map(dataset_id:str = "Saving-Willy/Happywhale-kaggle",
                    data_files:str = "data/train-00000-of-00001.parquet", 
                    dbg_show_extra:bool = False):
    '''
    render a map, with a selectable tileset, and show markers for each of the whale 
    observations
    
    '''
    # load/download data from huggingface dataset
    metadata = load_dataset(dataset_id, data_files=data_files)
    
    # make a pandas df that is compliant with folium/streamlit maps
    _df = pd.DataFrame({
        'lat': metadata["train"]["latitude"],
        'lon': metadata["train"]["longitude"],
        'species': metadata["train"]["predicted_class"],}
    )
    if dbg_show_extra:
        # add a few samples to visualise colours 
        _df.loc[len(_df)] = {'lat': 0, 'lon': 0, 'species': 'rough_toothed_dolphin'}
        _df.loc[len(_df)] = {'lat': -3, 'lon': 0, 'species': 'pygmy_killer_whale'}
        _df.loc[len(_df)] = {'lat': 45.7, 'lon': -2.6, 'species': 'humpback_whale'}

    ocean_loc = 0, 10
    selected_tile = st.selectbox("Choose a tile set", tile_sets, index=None, placeholder="Choose a tile set...", disabled=False)
    map_ = create_map(selected_tile, ocean_loc, zoom_start=2)

    folium.Marker(
        location=ocean_loc,
        popup="Atlantis",
        tooltip="Atlantis",
        icon=folium.Icon(color='blue', icon='info-sign')
    ).add_to(map_)
    
    for _, row in _df.iterrows():
        c = whale2color.get(row['species'], 'red')
        msg = f"[D] color for {row['species']} is {c}"
        m_logger.debug(msg) # depends on m_logger logging level (*not* the main st app's logger)
        #m_logger.info(msg)
        
        kw = {"prefix": "fa", "color": 'gray', "icon_color": c, "icon": "binoculars" } 
        folium.Marker(
            location=[row['lat'], row['lon']],
            popup=f"{row['species']} ",
            tooltip=row['species'],
            icon=folium.Icon(**kw)
        ).add_to(map_)
        #st.info(f"Added marker for {row['name']} {row['lat']} {row['lon']}")

    st_data = st_folium(map_, width=725)

    # workaround for correctly showing js components in tabs
    js_show_zeroheight_iframe(
        component_iframe_title="streamlit_folium.st_folium",
        height=800,
    )
    # this is just debug info -- 
    #st.info("[D]" + str(metadata.column_names))

    return st_data