app.py

import json
import os
import gradio as gr
import plotly.graph_objects as go

import sys

sys.path.append(os.path.join(os.path.dirname(__file__), "..", "SatSeg"))
from satseg.geo_tools import (
    shapefile_to_latlong,
    shapefile_to_grid_indices,
    points_to_shapefile,
    get_cached_grid_indices,
)

DATA_DIR = "data/"
MASK_PATH = os.path.join(DATA_DIR, "serp.shp")
HOSP_PATH = os.path.join(DATA_DIR, "hospitals.shp")
POLICE_PATH = os.path.join(DATA_DIR, "police.shp")
OUT_DIR = "out/"


def gr_generate_map(
    token: str,
    side_len: str,
    show_grid: bool,
    show_mask: bool,
    show_hospitals: bool = True,
    show_police: bool = True,
    region: str = "None",
):
    token = "pk.eyJ1IjoiZGlsaXRoIiwiYSI6ImNsaDQ3NXF3ZDAxdDMzZXMxeWJic2h1cDQifQ.DDczQCDfTgQEUt6pGvjUAg"
    side_len = float(side_len)

    scattermaps = []
    if show_grid:
        grid_path = MASK_PATH[: -len(".shp")] + f"-len={side_len}.shp"
        if not os.path.exists(grid_path):
            indices = shapefile_to_grid_indices(MASK_PATH, side_len)
            points_to_shapefile(indices, grid_path)
        else:
            indices, _ = get_cached_grid_indices(grid_path)
        box = go.Scattermapbox(
            lat=indices[:, 1],
            lon=indices[:, 0],
            mode="markers",
            marker=go.scattermapbox.Marker(size=6),
        )
        box.name = "Grids"
        scattermaps.append(box)
    if show_mask:
        contours = shapefile_to_latlong(MASK_PATH)
        for contour in contours:
            lons = contour[:, 0]
            lats = contour[:, 1]
            scattermaps.append(
                go.Scattermapbox(
                    fill="toself",
                    lat=lats,
                    lon=lons,
                    mode="markers",
                    marker=go.scattermapbox.Marker(size=6),
                )
            )

    if show_hospitals:
        indices, labels = get_cached_grid_indices(HOSP_PATH)
        box = go.Scattermapbox(
            lat=indices[:, 1],
            lon=indices[:, 0],
            mode="markers+text",
            text=labels,
            marker=go.scattermapbox.Marker(size=10),
        )
        box.name = "Hospitals"
        box.textfont.update({"color": "White"})
        scattermaps.append(box)

    if show_police:
        indices, labels = get_cached_grid_indices(POLICE_PATH)
        box = go.Scattermapbox(
            lat=indices[:, 1],
            lon=indices[:, 0],
            mode="markers+text",
            text=labels,
            marker=go.scattermapbox.Marker(size=10),
        )
        box.name = "Police Stations"
        box.textfont.update({"color": "White"})
        scattermaps.append(box)

    fig = go.Figure(scattermaps)

    center = (7.753769, 80.691730)
    if region == "Ussangoda":
        center = (6.0994295, 80.9860763)
    elif region == "Indikolapelessa":
        center = (6.3602253, 80.9371957)
    elif region == "Ginigalpelessa":
        center = (6.3846744, 80.8868755)
    elif region == "Yudhaganawa":
        center = (7.665643, 80.9529867)
    if token:
        fig.update_layout(
            mapbox=dict(
                style="satellite-streets",
                accesstoken=token,
                center=go.layout.mapbox.Center(lat=center[0], lon=center[1]),
                pitch=0,
                zoom=6 if region == "None" else 13,
            ),
            mapbox_layers=[
                {
                    # "below": "traces",
                    "sourcetype": "raster",
                    "sourceattribution": "United States Geological Survey",
                    "source": [
                        "https://basemap.nationalmap.gov/arcgis/rest/services/USGSImageryOnly/MapServer/tile/{z}/{y}/{x}"
                    ],
                }
            ],
        )
    else:
        fig.update_layout(
            mapbox_style="open-street-map",
            hovermode="closest",
            mapbox=dict(
                bearing=0,
                center=go.layout.mapbox.Center(lat=center[0], lon=center[1]),
                pitch=0,
                zoom=6 if region == "None" else 13,
            ),
        )

    return fig


with gr.Blocks() as demo:
    gr.Markdown("""# SatSeg""")

    with gr.Tab("Sampling"):
        grid_token = gr.Textbox(
            value="", label="Mapbox Token (https://account.mapbox.com/)"
        )
        grid_side_len = gr.Textbox(value="100", label="Sampling Gap (m)")

        grid_show_grid = gr.Checkbox(True, label="Show Grid")
        grid_show_mask = gr.Checkbox(False, label="Show Mask")
        grid_show_hosp = gr.Checkbox(True, label="Show Hospitals")
        grid_show_police = gr.Checkbox(True, label="Show Police Stations")

        grid_button = gr.Button("Generate Grid")

        grid_map = gr.Plot(label="Plot")

        grid_region = gr.Radio(
            label="Zoom to Region",
            choices=[
                "None",
                "Ussangoda",
                "Indikolapelessa",
                "Ginigalpelessa",
                "Yudhaganawa",
            ],
        )

    grid_button.click(
        gr_generate_map,
        inputs=[grid_token, grid_side_len, grid_show_grid, grid_show_mask],
        outputs=grid_map,
    )

    grid_region.change(
        gr_generate_map,
        inputs=[
            grid_token,
            grid_side_len,
            grid_show_grid,
            grid_show_mask,
            grid_show_hosp,
            grid_show_police,
            grid_region,
        ],
        outputs=grid_map,
    )

demo.queue(concurrency_count=10).launch(debug=True, share=True)