Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
-import torch
-import numpy as np
+import geopandas as gpd
+import leafmap.foliumap as leafmap
 from PIL import Image
 import rasterio
 from rasterio.windows import Window
@@ -11,11 +11,11 @@ import os
 import albumentations as albu
 import segmentation_models_pytorch as smp
 from albumentations.pytorch.transforms import ToTensorV2
-import geopandas as gpd
 from shapely.geometry import shape
 from shapely.ops import unary_union
 from rasterio.features import shapes
-
+import torch
+import numpy as np
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ENCODER = 'se_resnext50_32x4d'
@@ -25,7 +25,7 @@ ENCODER_WEIGHTS = 'imagenet'
 # Load and prepare the model
 @st.cache_resource
 def load_model():
-    model = torch.load('deeplabv3 v15.pth', map_location=DEVICE)
+    model = torch.load('deeplabv3+ v15.pth', map_location=DEVICE)
     model.eval().float()
     return model
 
@@ -76,7 +76,7 @@ def process_and_predict(image, model):
     return mask_image
 
 
-def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4,threshold=0.6):
+def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, threshold=0.6):
     tiles = []
 
     with rasterio.open(map_file) as src:
@@ -113,8 +113,8 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4,thresh
                         "width": tile_size,
                         "transform": rasterio.windows.transform(window, src.transform)
                     })
-                    tile_image = np.array(tile_image)
 
+                    tile_image = np.array(tile_image)
                     preprocessed_tile = preprocess(image=tile_image)['image']
                     batch_images.append(preprocessed_tile)
                     batch_metas.append(out_meta)
@@ -124,7 +124,7 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4,thresh
 
                 batch_tensor = torch.cat([img.unsqueeze(0).to(DEVICE) for img in batch_images], dim=0)
                 with torch.no_grad():
-                    batch_masks = model(batch_tensor.to(DEVICE))
+                    batch_masks = model(batch_tensor)
 
                 batch_masks = torch.sigmoid(batch_masks)
                 batch_masks = (batch_masks > threshold).float()
@@ -142,58 +142,76 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4,thresh
     return tiles
 
 
-
-
-
-
 def create_vector_mask(tiles, output_path):
     all_polygons = []
     for mask_array, meta in tiles:
-        # Ensure mask is binary
+
         mask_array = (mask_array > 0).astype(np.uint8)
 
-        # Get shapes from the mask
         mask_shapes = list(shapes(mask_array, mask=mask_array, transform=meta['transform']))
 
-        # Convert shapes to Shapely polygons
+        # to shapely polygons
         polygons = [shape(geom) for geom, value in mask_shapes if value == 1]
 
         all_polygons.extend(polygons)
-    # Perform union of all polygons
+    #union of all polygons
     union_polygon = unary_union(all_polygons)
-    # Create a GeoDataFrame
+    # create gdf
     gdf = gpd.GeoDataFrame({'geometry': [union_polygon]}, crs=meta['crs'])
     # Save to file
     gdf.to_file(output_path)
 
-    # Calculate area in square meters
+    #area in square meters
     area_m2 = gdf.to_crs(epsg=3857).area.sum()
-    # Convert to hectares
 
     return gdf, area_m2
 
 
+def display_map(shapefile_path, tif_path):
+    st.title("Map with Shape and TIFF Overlay")
+
+    mask = gpd.read_file(shapefile_path)
+
+    if mask.crs is None or mask.crs.to_string() != 'EPSG:3857':
+        mask = mask.to_crs('EPSG:3857')
+
+    bounds = mask.total_bounds  # [minx, miny, maxx, maxy]
+    center = [(bounds[1] + bounds[3]) / 2, (bounds[0] + bounds[2]) / 2]
+
+    m = leafmap.Map(
+        center=[center[1], center[0]],  # leafmap uses [latitude, longitude]
+        zoom=10,
+        crs='EPSG3857'
+    )
+
+    m.add_gdf(mask, layer_name="Shapefile Mask")
+
+    m.add_raster(tif_path, layer_name="Satellite Image", rgb=True, opacity=0.9)
+
+    m.to_streamlit()
+
+
 def main():
-    st.title("PV segmentor")
+    st.title("PV Segmentor")
 
     uploaded_file = st.file_uploader("Choose a TIF file", type="tif")
 
     if uploaded_file is not None:
         st.write("File uploaded successfully!")
 
-        threshold= st.slider(
-        'Select the value of the threshold',
-        min_value=0.1,
-        max_value=0.9,
-        value=0.5,
-        step=0.05
+        threshold = st.slider(
+            'Select the value of the threshold',
+            min_value=0.1,
+            max_value=0.9,
+            value=0.6,
+            step=0.05
         )
-        overlap= st.slider(
-        'Select the value of overlap',
-        min_value=50,
-        max_value=150,
-        value=100,
-        step=25
+        overlap = st.slider(
+            'Select the value of overlap',
+            min_value=50,
+            max_value=150,
+            value=100,
+            step=25
         )
         st.write('Selected threshold value:', threshold)
         st.write('Selected overlap value:', overlap)
@@ -205,7 +223,7 @@ def main():
                 f.write(uploaded_file.getbuffer())
 
             best_model.float()
-            tiles = extract_tiles("temp.tif", best_model, tile_size=512, overlap=overlap, batch_size=4,threshold=threshold)
+            tiles = extract_tiles("temp.tif", best_model, tile_size=512, overlap=overlap, batch_size=4, threshold=threshold)
 
             st.write("Processing complete!")
 
@@ -213,7 +231,7 @@ def main():
             result_gdf, area_m2 = create_vector_mask(tiles, output_path)
 
             st.write("Vector mask created successfully!")
-            st.write(f"Total area occupied by polygons: {area_m2:.4f} m^2")
+            st.write(f"Total area occupied by PV panels: {area_m2:.4f} m^2")
 
             # Offer the shapefile for download
             shp_files = [f for f in os.listdir() if
@@ -232,11 +250,12 @@ def main():
                     mime="application/zip"
                 )
 
-            # Clean up temporary files
-            os.remove("temp.tif")
-            for file in shp_files:
-                os.remove(file)
+            display_map("output_mask.shp", "temp.tif")
+
+            #os.remove("temp.tif")
+            #for file in shp_files:
+            #    os.remove(file)
 
 
 if __name__ == "__main__":
-    main()
+    main()