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.")