app.py

import os
import ee
import json
import geemap
import geemap.foliumap as gee_folium
import leafmap.foliumap as leaf_folium
import streamlit as st
from pandas import to_datetime, read_csv, merge, date_range, DateOffset
from geopandas import read_file
from shapely.ops import transform
from functools import reduce
import plotly.express as px

st.set_page_config(layout="wide")

############################################
# One time setup
############################################


def initialize_ee():
    credentials_path = os.path.expanduser("~/.config/earthengine/credentials")
    if os.path.exists(credentials_path):
        pass  # Earth Engine credentials already exist
    elif "EE" in os.environ:  # write the credentials to the file
        ee_credentials = os.environ.get("EE")
        os.makedirs(os.path.dirname(credentials_path), exist_ok=True)
        with open(credentials_path, "w") as f:
            f.write(ee_credentials)
    else:
        raise ValueError(
            f"Earth Engine credentials not found at {credentials_path} or in the environment variable 'EE'"
        )

    ee.Initialize()


if "ee_initialized" not in st.session_state:
    initialize_ee()
    st.session_state.ee_initialized = True
    
if "wayback_mapping" not in st.session_state:
    with open("wayback_imagery.json") as f:
        st.session_state.wayback_mapping = json.load(f)


############################################
# Functions
############################################
def shape_3d_to_2d(shape):
    if shape.has_z:
        return transform(lambda x, y, z: (x, y), shape)
    else:
        return shape


def preprocess_gdf(gdf):
    gdf = gdf.to_crs(epsg=4326)
    gdf = gdf[['Name', 'geometry']]
    gdf["geometry"] = gdf["geometry"].apply(shape_3d_to_2d)
    return gdf


def calculate_ndvi(image, nir_band, red_band):
    nir = image.select(nir_band)
    red = image.select(red_band)
    ndvi = (nir.subtract(red)).divide(nir.add(red)).rename("NDVI")
    return image.addBands(ndvi)


def postprocess_df(df, name):
    df = df.T
    df = df.reset_index()
    ndvi_df = df[df["index"].str.contains("NDVI")]
    ndvi_df["index"] = to_datetime(ndvi_df["index"], format="%Y-%m_NDVI")
    ndvi_df = ndvi_df.rename(columns={"index": "Date", 0: name})
    
    cloud_mask_probability = df[df["index"].str.contains("MSK_CLDPRB")]
    cloud_mask_probability["index"] = to_datetime(cloud_mask_probability["index"], format="%Y-%m_MSK_CLDPRB")
    cloud_mask_probability = cloud_mask_probability.rename(columns={"index": "Date", 0: f"{name}_cloud_proba"})
    # normalize
    cloud_mask_probability[f"{name}_cloud_proba"] = cloud_mask_probability[f"{name}_cloud_proba"] / 100
    df = merge(ndvi_df, cloud_mask_probability, on="Date", how="outer")
    return df


def write_info(info):
    st.write(f"<span style='color:#00FF00;'>{info}</span>", unsafe_allow_html=True)


############################################
# App
############################################

# Title
# make title in center
st.markdown(
    f"""
    <h1 style="text-align: center;">Mean NDVI Calculator</h1>
    """,
    unsafe_allow_html=True,
)

# Input: Date and Cloud Cover
col = st.columns(2)
start_date = col[0].date_input("Start Date", value=to_datetime("2021-01-01"))
end_date = col[1].date_input("End Date", value=to_datetime("2021-07-31"))
start_date = start_date.strftime("%Y-%m")
end_date = end_date.strftime("%Y-%m")
# max_cloud_cover = st.number_input("Max Cloud Cover (in percentage)", value=5)

# Input: GeoJSON/KML file
uploaded_file = st.file_uploader("Upload KML/GeoJSON file", type=["geojson", "kml"])
if uploaded_file is None:
    st.stop()

file_name = uploaded_file.name
gdf = read_file(uploaded_file)
gdf = preprocess_gdf(gdf)

selected_shape = st.selectbox("Select the geometry", gdf.Name.values)
if selected_shape is None:
    st.stop()

selected_shape = gdf[gdf.Name == selected_shape]
ee_object = geemap.gdf_to_ee(selected_shape)
write_info(f"Type of Geometry: {selected_shape.geometry.type.values[0]}")

st.write("Select the satellite sources:")
satellites = {
    "LANDSAT/LC08/C02/T1_TOA": {
        "selected": st.checkbox("LANDSAT/LC08/C02/T1_TOA", value=True),
        "nir_band": "B5",
        "red_band": "B4",
        "scale": 30,
    },
    "COPERNICUS/S2_SR_HARMONIZED": {
        "selected": st.checkbox("COPERNICUS/S2_SR_HARMONIZED", value=True),
        "nir_band": "B8",
        "red_band": "B4",
        "scale": 10,
    },
}

submit = st.button("Submit", use_container_width=True)
if submit:
    if not any(satellites.values()):
        st.error("Please select at least one satellite source")
        st.stop()

    # Create month range
    dates = date_range(start_date, end_date, freq="MS").strftime("%Y-%m-%d").tolist()
    write_info(
        f"Start Date (inclusive): {start_date}, End Date (exclusive): {end_date}"
    )

    df_list = []
    collections = {}
    for satellite, attrs in satellites.items():
        if not attrs["selected"]:
            continue
        collection = ee.ImageCollection(satellite)
        collection = collection.filterBounds(ee_object)
        if satellite == "COPERNICUS/S2_SR_HARMONIZED":
            collection = collection.select([attrs["red_band"], attrs["nir_band"], "MSK_CLDPRB"])
        else:
            collection = collection.select([attrs["red_band"], attrs["nir_band"]])
        # collection = collection.filter(ee.Filter.lt(attrs["cloud_cover_var"], max_cloud_cover))
        collection = collection.filterDate(start_date, end_date)
        collection = collection.map(
            lambda image: calculate_ndvi(
                image, nir_band=attrs["nir_band"], red_band=attrs["red_band"]
            )
        )
        write_info(f"Number of images in {satellite}: {collection.size().getInfo()}")
        
        progress_bar = st.progress(0)

        def monthly_quality_mosaic(start, end, i):
            progress_bar.progress((i + 1) / (len(dates) - 1))
            collection_filtered = collection.filterDate(start, end)
            size = collection_filtered.size().getInfo()
            if size == 0:
                return None
            mosaic = collection_filtered.qualityMosaic("NDVI")
            month = to_datetime(start).strftime("%Y-%m")
            print(f"Processing {month} with {size} images")
            return mosaic.set("system:index", f"{month}")
            
        collection = [monthly_quality_mosaic(start, end, i) for i, (start, end) in enumerate(zip(dates[:-1], dates[1:]))]
        collection = list(filter(None, collection))
        collection = ee.ImageCollection(collection)
        collections[satellite] = collection

        save_name = satellite.replace("/", "_")
        geemap.zonal_stats(
            collection,#.select(["NDVI"]),
            ee_object,
            f"/tmp/{save_name}.csv",
            stat_type="mean",
            scale=attrs["scale"],
        )
        df = read_csv(f"/tmp/{save_name}.csv")
        df = postprocess_df(df, name=satellite)
        df_list.append(df)

    df = reduce(lambda left, right: merge(left, right, on="Date", how="outer"), df_list)
    df = df.sort_values("Date")
    
    # drop rows with all NaN values
    df = df.dropna(how="all")
    # drop columns with all NaN values
    df = df.dropna(axis=1, how="all")
    df = df.reset_index(drop=True)
    
    st.session_state.df = df
    st.session_state.collections = collections

if "df" in st.session_state:
    df = st.session_state.df
    collections = st.session_state.collections
    st.write(df.applymap(lambda x: f"{x:.2f}" if isinstance(x, float) else x))

    fig = px.line(df, x="Date", y=df.columns[1:], title='Mean NDVI', markers=True)
    fig.update_yaxes(range=[0, 1])
    st.plotly_chart(fig)
    
    st.subheader("Visual Inspection")
    write_info(f"Centroid of the selected geometry: {selected_shape.geometry.centroid.values[0]}")
    cols = st.columns(2)
    with cols[0]:
        start_date = st.selectbox("Start Date", df.Date, index=0)
        start_date_index = df[df.Date == start_date].index[0].item()
    with cols[1]:
        end_date = st.selectbox("End Date", df.Date, index=len(df.Date) - 1)
        end_date_index = df[df.Date == end_date].index[0].item()
    
    for imagery in satellites:
        collection = collections[imagery]
        for col, date in zip(cols, [start_date, end_date]):
            date_index = df[df.Date == date].index[0].item()
            image = ee.Image(collections[imagery].toList(collection.size()).get(date_index))
            layer = gee_folium.ee_tile_layer(image, {"bands": ["NDVI"], "min": 0, "max": 1}, f"{imagery}_{date}")

            with col:
                m = leaf_folium.Map()
                m.add_layer(layer)
                m.add_gdf(selected_shape, layer_name="Selected Geometry")
                colors = ["#000000", "#FFFFFF"]
                m.add_colorbar(colors, vmin=0, vmax=1)
                st.write(f"{imagery} - {date}")
                m.to_streamlit()
    
    for col, date in zip(cols, [start_date, end_date]):
        esri_date = min(st.session_state.wayback_mapping.keys(), key=lambda x: abs(to_datetime(x) - date))
        with col:
            m = leaf_folium.Map()
            m.add_tile_layer(st.session_state.wayback_mapping[esri_date], name=f"Esri Wayback Imagery - {esri_date}", attribution="Esri")
            m.add_gdf(selected_shape, layer_name="Selected Geometry")
            st.write(f"Esri Wayback Imagery - {esri_date} (Closest to {date})")
            m.to_streamlit()