Spaces:
Sleeping
Sleeping
File size: 3,717 Bytes
a2da3c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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)
|