app.py

import itertools
from datetime import datetime
from pathlib import Path

import folium
import streamlit as st
from loguru import logger as log
from matplotlib import colors
from streamlit_folium import st_folium
from utils.folium import (
    draw_gdf,
    get_clean_rendering_container,
    get_map_crop,
    get_parquet,
)

COLOR_MAP = {
    "Grassland": "lime",
    "Cereals": "saddlebrown",
    "Other": "white",
    "Flowers": "cyan",
    "Vegetables": "purple",
    "Fruit trees": "darkgreen",
}


id_class = {(i + 1): k for i, k in enumerate(COLOR_MAP.keys())}
color_map_hex = {k: colors.to_hex(v) for k, v in COLOR_MAP.items()}
id_map_hex = {(i + 1): v for i, v in enumerate(color_map_hex.values())}
color_map_hex_r = {v: k for k, v in color_map_hex.items()}
merged_list = list(itertools.chain(*list(color_map_hex_r.items())))


def style_call(feat):
    di = {
        "fillColor": id_map_hex[feat["properties"]["class"]],
        "fillOpacity": 0.7,
        "color": id_map_hex[feat["properties"]["class"]],
        "border-width": "thin",
        "border-color": "#ffffff",
        "weight": 1.2,
    }
    return di


# Page configs
st.set_page_config(page_title="AIDA", page_icon="🌐", layout="wide")

hide_button_style = """
    <style>
    div.stButton > button:first-child {
        display: none;
    }
    </style>
"""
st.markdown(hide_button_style, unsafe_allow_html=True)

# base paths
base_path = Path("data")
inference_path = base_path / "inference"

# list examples and predictions
inference_paths = sorted(list(inference_path.iterdir()))

example_names = [p.stem for p in inference_paths]
inference_dict = {p.stem: p for p in inference_paths}

# Create selection menu
container_predictions = st.sidebar.container(border=True)
with container_predictions:
    prediction_selectbox = st.selectbox(
        "Select an example",
        options=example_names,
        index=None,
        key="selectbox_pred",
    )
    is_prediction_selected = prediction_selectbox is not None
    if is_prediction_selected:
        try:
            # add loading of the selected parquet prediction
            chosen_tile_path = inference_dict[prediction_selectbox]
            with open(chosen_tile_path, "rb") as f:
                prediction_file = f.read()
        except:
            st.warning("File not found")
            chosen_tile_path = None
            prediction_file = None
    else:
        prediction_file = None
        chosen_tile_path = None

container = get_clean_rendering_container(prediction_selectbox)


# Stange Hack to always update the map height I guess
def change_key():
    st.session_state["key_map"] = str(datetime.now())


height_value = st.sidebar.radio(
    "Map height",
    options=[800, 500],
    horizontal=True,
    key="height_map",
    on_change=change_key,
)
container = get_clean_rendering_container(height_value)

# draw map
interactive_map = get_map_crop()

if prediction_selectbox is not None:
    # draw prediction
    draw_gdf(interactive_map, chosen_tile_path, "Prediction", style_call, id_class)


with container.form(key="form1", clear_on_submit=True):
    folium.LayerControl().add_to(interactive_map)
    output_map = st_folium(
        interactive_map,
        width=None,
        height=height_value,
        returned_objects=["all_drawings"],
        key=st.session_state.get("key_map", "key_map"),
    )
    submit = st.form_submit_button("Recenter map")

# Update messages


for name, path in inference_dict.items():
    if path not in st.session_state.keys():
        log.info(f"Loading parquet {name}")
        st.session_state[path] = get_parquet(path, id_class)