Sync with https://github.com/mozilla-ai/osm-ai-helper

Dockerfile

@@ -0,0 +1,26 @@
+FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04
+
+RUN apt-get update && apt-get install --no-install-recommends -y \
+  espeak-ng \
+  build-essential \
+  python3.10 \
+  python3.10-dev \
+  python3-pip \
+  git \
+  && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 user
+
+USER user
+
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+RUN pip3 install git+https://github.com/mozilla-ai/osm-ai-helper.git
+
+COPY --chown=user . $HOME/app
+
+EXPOSE 8501
+ENTRYPOINT ["streamlit", "run", "app.py", "--server.enableXsrfProtection", "false"]

app.py

@@ -0,0 +1,119 @@
+import os
+from pathlib import Path
+from shutil import move
+
+import folium
+import streamlit as st
+from loguru import logger
+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
+
+
+class StreamlitHandler:
+    def __init__(self):
+        self.text_widget = st.empty()
+        self.log_messages = ""
+
+    def write(self, message):
+        self.log_messages += message
+        self.text_widget.code(self.log_messages)
+        return
+
+
+@st.fragment
+def inference(lat_lon, margin):
+    hf_hub_download(
+        "daavoo/yolo-osm-pool-detector",
+        filename="model.pt",
+        repo_type="model",
+        local_dir="models",
+    )
+    output_path, *_ = run_inference(
+        model_file="models/model.pt",
+        output_dir="results",
+        lat_lon=lat_lon,
+        margin=margin,
+    )
+    return output_path
+
+
+@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_column_width=True)
+    with col2:
+        st.image(painted_image, caption="Painted Image", use_column_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.success(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 our OpenStreetMap account."
+        "You can check the [Colab Notebook](ttps://colab.research.google.com/github/mozilla-ai/osm-ai-helper/blob/main/demo/run_inference.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"
+        upload_osm(
+            results_dir=output_path / "keep",
+            client_id=os.environ["OSM_CLIENT_ID"],
+            client_secret=os.environ["OSM_CLIENT_SECRET"],
+            comment=comment,
+        )
+
+
+st.title("Open Street Map AI Helper")
+
+st.divider()
+
+st.subheader("Click on the map to select a latitude and longitude")
+
+m = folium.Map(location=[42.2489, -8.5117], zoom_start=11, tiles="OpenStreetMap")
+
+st_data = st_folium(m, width=725)
+
+if st_data.get("last_clicked"):
+    lat = st_data["last_clicked"]["lat"]
+    lon = st_data["last_clicked"]["lng"]
+    st.write(f"Last Clicked: {lat}, {lon}")
+
+    if st.button("Run Inference"):
+        streamlit_handler = StreamlitHandler()
+        logger.add(streamlit_handler, format="<level>{message}</level>")
+
+        output_path = inference(lat_lon=(lat, lon), margin=1)
+
+        for new in Path(output_path).glob("*.json"):
+            handle_polygon(new)
+
+        upload_results(output_path)