osm-ai-helper / app.py
github-actions[bot]
Sync with https://github.com/mozilla-ai/osm-ai-helper
cd52496
import os
from pathlib import Path
from shutil import move
import folium
import streamlit as st
from branca.element import MacroElement
from jinja2 import Template
from huggingface_hub import hf_hub_download
from PIL import Image
from streamlit_folium import st_folium
from osm_ai_helper.run_inference import run_inference
from osm_ai_helper.upload_osm import upload_osm
@st.fragment
def show_map():
class LatLngPopup(MacroElement):
_template = Template(
"""
{% macro script(this, kwargs) %}
var {{this.get_name()}} = L.popup();
function latLngPop(e) {
{{this.get_name()}}
.setLatLng(e.latlng)
.setContent(e.latlng.lat.toFixed(4) + ", " + e.latlng.lng.toFixed(4))
.openOn({{this._parent.get_name()}});
}
{{this._parent.get_name()}}.on('click', latLngPop);
{% endmacro %}
"""
)
def __init__(self):
super().__init__()
self._name = "LatLngPopup"
m = folium.Map(location=[42.8075, -8.1519], zoom_start=8, tiles="OpenStreetMap")
m.add_child(LatLngPopup())
st_folium(m, height=400, width=800)
@st.fragment
def inference(lat_lon):
with st.spinner("Downloading model..."):
hf_hub_download(
"mozilla-ai/swimming-pool-detector",
filename="model.pt",
repo_type="model",
local_dir="models",
)
with st.spinner("Downloading image and Running inference..."):
output_path, existing, new, missed = run_inference(
yolo_model_file="models/model.pt",
output_dir="/tmp/results",
lat_lon=lat_lon,
margin=2,
save_full_images=False,
batch_size=64,
)
return output_path, existing, new
@st.fragment
def handle_polygon(polygon):
raw_image = Image.open(polygon.with_suffix(".png"))
painted_image = Image.open(f"{polygon.parent}/{polygon.stem}_painted.png")
st.subheader(f"Reviewing: {polygon.name}")
col1, col2 = st.columns(2)
with col1:
st.image(raw_image, caption="Raw Image", use_container_width=True)
with col2:
st.image(painted_image, caption="Painted Image", use_container_width=True)
if st.button("Keep Polygon", key=f"keep_{polygon}"):
keep_folder = polygon.parent / "keep"
keep_folder.mkdir(parents=True, exist_ok=True)
move(polygon, keep_folder / polygon.name)
st.success(f"Polygon moved to {keep_folder}")
elif st.button("Discard Polygon", key=f"discard_{polygon.stem}"):
discard_folder = polygon.parent / "discard"
discard_folder.mkdir(parents=True, exist_ok=True)
move(polygon, discard_folder / polygon.name)
st.warning(f"Polygon moved to {discard_folder}")
@st.fragment
def upload_results(output_path):
st.divider()
st.header("Upload all polygons in `keep`")
st.markdown(
"The results will be uploaded using the [osm-ai-helper](https://www.openstreetmap.org/user/osm-ai-helper) profile."
)
st.markdown(
"You can check the [Colab Notebook](ttps://colab.research.google.com/github/mozilla-ai/osm-ai-helper/blob/main/demo/run_inference_point.ipynb)"
" and the [Authorization Guide](https://mozilla-ai.github.io/osm-ai-helper/authorization)"
" to contribute with your own OpenStreetMap account."
)
contributor = st.text_input("(Optional) Indicate your name for attribution")
if st.button("Upload all polygons in `keep`"):
if contributor:
comment = f"Add Swimming Pools. Contributed by {contributor}"
else:
comment = "Add Swimming Pools"
changeset = upload_osm(
results_dir=output_path / "keep",
client_id=os.environ["OSM_CLIENT_ID"],
client_secret=os.environ["OSM_CLIENT_SECRET"],
comment=comment,
)
st.success(
f"Changeset created: https://www.openstreetmap.org/changeset/{changeset}"
)
st.title("OpenStreetMap AI Helper")
st.markdown(
"""
This demo was created with the repo [mozilla-ai/osm-ai-helper](https://github.com/mozilla-ai/osm-ai-helper).
It uses the model [mozilla-ai/swimming-pool-detector](https://huggingface.co/mozilla-ai/swimming-pool-detector).
You can check the [Create Dataset](https://colab.research.google.com/github/mozilla-ai//osm-ai-helper/blob/main/demo/create_dataset.ipyn)
and [Finetune Model](https://colab.research.google.com/github/mozilla-ai//osm-ai-helper/blob/main/demo/finetune_model.ipynb) notebooks to learn how to train your own model.
"""
)
st.divider()
st.subheader("Click on the map to select a latitude and longitude.")
st.markdown(
"""
The model will try to find swimming pools around this location.
Note that this model was trained with data from [Galicia](https://nominatim.openstreetmap.org/ui/details.html?osmtype=R&osmid=349036&class=boundary),
so it might fail to generalize to significantly different places.
"""
)
show_map()
lat_lon = st.text_input("Paste the copied (latitude, longitude)")
if st.button("Run Inference") and lat_lon:
lat, lon = lat_lon.split(",")
output_path, existing, new = inference(
lat_lon=(float(lat.strip()), float(lon.strip()))
)
st.info(f"Found {len(existing)} swimming pools already in OpenStreetMaps.")
if new:
st.divider()
st.header("Review `new` swimming pools")
st.markdown(
"Every `new` swimming pool will be displayed at the center of the image in `yellow`."
)
st.markdown(
"Swimming pools in other colors are those already existing in OpenStreetMap and they just "
"indicate whether the model has found them (`green`) or missed them (`red`)."
)
for new in Path(output_path).glob("*.json"):
handle_polygon(new)
upload_results(output_path)
else:
st.warning("No `new` swimming pools were found. Try a different location.")