Spaces:
Running
on
T4
Running
on
T4
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 | |
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) | |
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 | |
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}") | |
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.") | |