Update pages/12_🌲_VertXtractor.py

pages/12_🌲_VertXtractor.py

@@ -6,22 +6,14 @@ from folium.plugins import Draw
 import requests
 from shapely.geometry import box
 import json
+import io
+import rasterio
+from rasterio.io import MemoryFile
+import numpy as np
 
 st.set_page_config(layout="wide", page_title="Extracteur de données géospatiales")
 
-def get_swisstopo_collections():
-    url = "https://data.geo.admin.ch/api/stac/v0.9/collections"
-    response = requests.get(url)
-    if response.status_code == 200:
-        collections = response.json()['collections']
-        return {c['title']: {"id": c['id'], "source": "swisstopo", "type": c.get('type', 'unknown')} for c in collections}
-    else:
-        st.error("Impossible de récupérer les collections Swisstopo")
-        return {}
-
-SWISSTOPO_ALL = get_swisstopo_collections()
-
-SWISSTOPO_SELECTED = {
+SWISSTOPO_LAYERS = {
     "SWISSIMAGE 10 cm": {"id": "ch.swisstopo.swissimage-dop10", "source": "swisstopo", "type": "image"},
     "Carte nationale 1:25'000": {"id": "ch.swisstopo.pixelkarte-farbe-pk25.noscale", "source": "swisstopo", "type": "image"},
     "MNT": {"id": "ch.swisstopo.swissalti3d", "source": "swisstopo", "type": "raster"},
@@ -50,64 +42,114 @@ def create_geojson_from_box(bbox):
     return json.loads(gdf.to_json()), json.loads(gdf.to_crs("EPSG:2056").to_json())
 
 @st.cache_data
-def get_download_url(bbox, layer_info):
-    urls = {"swisstopo": None, "esri": None}
-    
-    # Swisstopo
-    if layer_info["source"] == "swisstopo":
-        api_url = f"https://data.geo.admin.ch/api/stac/v0.9/collections/{layer_info['id']}/items"
-        params = {"bbox": ",".join(map(str, bbox)), "limit": 1}
+def extract_swisstopo_data(bbox, layer_id, layer_type):
+    api_url = f"https://data.geo.admin.ch/api/stac/v0.9/collections/{layer_id}/items"
+    params = {"bbox": ",".join(map(str, bbox)), "limit": 1}
+    try:
         response = requests.get(api_url, params=params)
+        response.raise_for_status()
         data = response.json()
-        if 'features' in data and data['features']:
+        if data['features']:
             feature = data['features'][0]
-            assets = feature.get('assets', {})
-            if layer_info['id'] == 'ch.swisstopo.swissimage-dop10':
-                data_key = 'rgb' if 'rgb' in assets else 'data' if 'data' in assets else None
-            else:
-                data_key = 'data' if 'data' in assets else next(iter(assets)) if assets else None
-            
-            if data_key and data_key in assets:
-                urls["swisstopo"] = assets[data_key].get('href')
-    
-    # ESRI
-    if layer_info["id"] in [layer["id"] for layer in ESRI_LAYERS.values()]:
-        service_url = f"https://services.arcgisonline.com/arcgis/rest/services/{layer_info['id']}/MapServer/export"
-        params = {
-            "bbox": f"{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}",
-            "bboxSR": 4326,
-            "size": "1000,1000",
-            "format": "png",
-            "f": "html",
-            "imageSR": 4326,
-            "transparent": "true"
-        }
-        urls["esri"] = f"{service_url}?{'&'.join([f'{k}={v}' for k, v in params.items()])}"
-    
-    return urls
+            asset_keys = feature['assets'].keys()
+            data_key = 'rgb' if 'rgb' in asset_keys else 'data' if 'data' in asset_keys else next(iter(asset_keys))
+            data_url = feature['assets'][data_key]['href']
+            return requests.get(data_url).content
+    except requests.RequestException as e:
+        st.error(f"Erreur Swisstopo: {str(e)}")
+    return None
+
+@st.cache_data
+def extract_esri_data(bbox, layer_id, layer_type):
+    service_url = f"https://services.arcgisonline.com/arcgis/rest/services/{layer_id}/MapServer/export"
+    params = {
+        "bbox": f"{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}",
+        "bboxSR": 4326,
+        "size": "1000,1000",
+        "format": "tiff",
+        "f": "image"
+    }
+    try:
+        response = requests.get(service_url, params=params)
+        response.raise_for_status()
+        return response.content
+    except requests.RequestException as e:
+        st.error(f"Erreur ESRI: {str(e)}")
+    return None
+
+def process_and_display_data(data, layer_name, layer_type):
+    if data is None:
+        st.warning(f"Pas de données pour {layer_name}")
+        return
+    try:
+        with MemoryFile(data) as memfile:
+            with memfile.open() as src:
+                image = src.read()
+                if image.shape[0] == 3:
+                    image = np.transpose(image, (1, 2, 0))
+                elif image.shape[0] == 1:
+                    image = image[0]
+                
+                st.image(image, caption=layer_name, use_column_width=True)
+                
+                output = io.BytesIO()
+                with MemoryFile(output) as memfile:
+                    with memfile.open(**src.profile) as dst:
+                        dst.write(src.read())
+                
+                st.download_button(
+                    label=f"Télécharger {layer_name} (GeoTIFF)",
+                    data=output.getvalue(),
+                    file_name=f"{layer_name.lower().replace(' ', '_')}.tif",
+                    mime="image/tiff"
+                )
+    except Exception as e:
+        st.error(f"Erreur de traitement pour {layer_name}: {str(e)}")
 
 def main():
-    st.title("Extracteur de données géospatiales")
+    st.title("Extracteur de données Swisstopo et ESRI")
 
-    col1, col2 = st.columns([1, 2])
+    if 'selected_layers' not in st.session_state:
+        st.session_state.selected_layers = set()
+
+    col1, col2 = st.columns([1, 3])
 
     with col1:
-        st.subheader("1. Sélectionner les couches")
-        source = st.radio("Source des données", ["Swisstopo (Sélection)", "Swisstopo (Toutes)", "ESRI"])
+        st.subheader("Sélection des couches")
+        source = st.radio("Source des données", ["Swisstopo", "ESRI"])
         
-        if source == "Swisstopo (Sélection)":
-            layers = SWISSTOPO_SELECTED
-        elif source == "Swisstopo (Toutes)":
-            layers = SWISSTOPO_ALL
-        else:
-            layers = ESRI_LAYERS
-
-        selected_layers = st.multiselect("Couches à extraire", list(layers.keys()))
+        layers = SWISSTOPO_LAYERS if source == "Swisstopo" else ESRI_LAYERS
+        
+        for layer_name in layers:
+            if st.checkbox(layer_name, value=layer_name in st.session_state.selected_layers):
+                st.session_state.selected_layers.add(layer_name)
+            else:
+                st.session_state.selected_layers.discard(layer_name)
+        
+        st.markdown("---")
+        st.markdown("## À propos\nExtracteur de données géospatiales Swisstopo et ESRI.\nDéveloppé par Vertdure")
 
-        st.subheader("2. Définir la zone d'intérêt")
+    with col2:
+        st.subheader("1. Définir la zone d'intérêt")
         method = st.radio("Méthode", ["Dessiner", "GeoJSON"], horizontal=True)
-
-        if method == "GeoJSON":
+        if method == "Dessiner":
+            if 'bbox' not in st.session_state:
+                map_data = draw_box_on_map()
+                if map_data['last_active_drawing']:
+                    coords = map_data['last_active_drawing']['geometry']['coordinates'][0]
+                    bbox = box(min(c[0] for c in coords), min(c[1] for c in coords),
+                               max(c[0] for c in coords), max(c[1] for c in coords))
+                    st.session_state['bbox'] = bbox
+                    st.session_state['geojson_wgs84'], st.session_state['geojson_swiss'] = create_geojson_from_box(bbox)
+                    st.success("Zone sélectionnée!")
+            else:
+                st.success("Zone déjà sélectionnée. Cliquez sur 'Effacer la zone' pour redessiner.")
+                if st.button("Effacer la zone"):
+                    del st.session_state['bbox']
+                    del st.session_state['geojson_wgs84']
+                    del st.session_state['geojson_swiss']
+                    st.experimental_rerun()
+        else:
             uploaded_file = st.file_uploader("GeoJSON", type="geojson")
             if uploaded_file:
                 gdf = gpd.read_file(uploaded_file)
@@ -115,40 +157,20 @@ def main():
                 st.session_state['geojson_swiss'] = json.loads(gdf.to_crs(2056).to_json())
                 st.success("GeoJSON chargé!")
 
-        st.subheader("3. Télécharger les données")
-        if st.button("Obtenir les liens de téléchargement"):
+        if st.button("Extraire les données"):
             if 'geojson_wgs84' not in st.session_state or 'geojson_swiss' not in st.session_state:
                 st.error("Définissez d'abord une zone d'intérêt.")
-            elif not selected_layers:
+            elif not st.session_state.selected_layers:
                 st.error("Sélectionnez au moins une couche à extraire.")
             else:
-                for layer_name in selected_layers:
-                    layer_info = layers[layer_name]
-                    bbox = gpd.GeoDataFrame.from_features(st.session_state['geojson_swiss' if layer_info["source"] == "swisstopo" else 'geojson_wgs84']).total_bounds
-                    st.write(f"Traitement de la couche : {layer_name}")
-                    st.write(f"Bbox : {bbox}")
-                    try:
-                        download_urls = get_download_url(bbox, layer_info)
-                        if download_urls["swisstopo"]:
-                            st.markdown(f"[Télécharger {layer_name} (Swisstopo)]({download_urls['swisstopo']})")
-                        if download_urls["esri"]:
-                            st.markdown(f"[Télécharger {layer_name} (ESRI)]({download_urls['esri']})")
-                        if not download_urls["swisstopo"] and not download_urls["esri"]:
-                            st.warning(f"Aucun lien de téléchargement disponible pour {layer_name}")
-                    except Exception as e:
-                        st.error(f"Erreur lors de la récupération du lien pour {layer_name}: {str(e)}")
-
-    with col2:
-        if method == "Dessiner":
-            map_data = draw_box_on_map()
-            if map_data['last_active_drawing']:
-                coords = map_data['last_active_drawing']['geometry']['coordinates'][0]
-                bbox = box(min(c[0] for c in coords), min(c[1] for c in coords),
-                           max(c[0] for c in coords), max(c[1] for c in coords))
-                st.session_state['geojson_wgs84'], st.session_state['geojson_swiss'] = create_geojson_from_box(bbox)
-                st.success("Zone sélectionnée!")
-
-    st.sidebar.markdown("## À propos\nExtracteur de données géospatiales Swisstopo et ESRI.\nDéveloppé par Vertdure")
+                with st.spinner("Extraction des données en cours..."):
+                    for layer_name in st.session_state.selected_layers:
+                        layer_info = SWISSTOPO_LAYERS.get(layer_name) or ESRI_LAYERS.get(layer_name)
+                        bbox = gpd.GeoDataFrame.from_features(st.session_state['geojson_swiss' if layer_info["source"] == "swisstopo" else 'geojson_wgs84']).total_bounds
+                        st.info(f"Traitement de la couche : {layer_name}")
+                        data = extract_swisstopo_data(bbox, layer_info["id"], layer_info["type"]) if layer_info["source"] == "swisstopo" else extract_esri_data(bbox, layer_info["id"], layer_info["type"])
+                        process_and_display_data(data, layer_name, layer_info["type"])
+                st.success("Extraction terminée !")
 
 if __name__ == "__main__":
     main()