Update app.py

app.py

@@ -22,10 +22,10 @@ ENCODER = 'se_resnext50_32x4d'
 ENCODER_WEIGHTS = 'imagenet'
 
 
-# Load and prepare the model
+#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
 
@@ -41,9 +41,9 @@ def to_tensor(x, **kwargs):
 preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
 
 
-def get_preprocessing():
+def get_preprocessing(tile_size):
     _transform = [
-        albu.Resize(512, 512),
+        albu.PadIfNeeded(min_height=tile_size, min_width=tile_size, always_apply=True),
         albu.Lambda(image=preprocessing_fn),
         albu.Lambda(image=to_tensor, mask=to_tensor),
         ToTensorV2(),
@@ -51,32 +51,13 @@ def get_preprocessing():
     return albu.Compose(_transform)
 
 
-preprocess = get_preprocessing()
 
 
-@torch.no_grad()
-def process_and_predict(image, model):
-    if isinstance(image, Image.Image):
-        image = np.array(image)
 
-    if image.ndim == 2:
-        image = np.stack([image] * 3, axis=-1)
-    elif image.shape[2] == 4:
-        image = image[:, :, :3]
-
-    preprocessed = preprocess(image=image)['image']
-    input_tensor = preprocessed.unsqueeze(0).to(DEVICE)
-
-    mask = model(input_tensor)
-    mask = torch.sigmoid(mask)
-    mask = (mask > 0.6).float()
-
-    mask_image = Image.fromarray((mask.squeeze().cpu().numpy() * 255).astype(np.uint8))
-
-    return mask_image
+def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, threshold=0.6):
 
+    preprocess = get_preprocessing(tile_size)
 
-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:
@@ -145,23 +126,24 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, thres
 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']))
 
-        # to shapely polygons
+        # Convert shapes to Shapely polygons
         polygons = [shape(geom) for geom, value in mask_shapes if value == 1]
 
         all_polygons.extend(polygons)
-    #union of all polygons
+    # Perform union of all polygons
     union_polygon = unary_union(all_polygons)
-    # create gdf
+    # Create a GeoDataFrame
     gdf = gpd.GeoDataFrame({'geometry': [union_polygon]}, crs=meta['crs'])
     # Save to file
     gdf.to_file(output_path)
 
-    #area in square meters
+    # Calculate area in square meters
     area_m2 = gdf.to_crs(epsg=3857).area.sum()
 
     return gdf, area_m2
@@ -170,24 +152,31 @@ def create_vector_mask(tiles, output_path):
 def display_map(shapefile_path, tif_path):
     st.title("Map with Shape and TIFF Overlay")
 
+    # Load the shapefile
     mask = gpd.read_file(shapefile_path)
 
+    # Check and reproject the mask to EPSG:3857 if needed
     if mask.crs is None or mask.crs.to_string() != 'EPSG:3857':
         mask = mask.to_crs('EPSG:3857')
 
+    # Get the bounds of the shapefile to center the map
     bounds = mask.total_bounds  # [minx, miny, maxx, maxy]
     center = [(bounds[1] + bounds[3]) / 2, (bounds[0] + bounds[2]) / 2]
 
+    # Create a leafmap centered on the shapefile bounds
     m = leafmap.Map(
         center=[center[1], center[0]],  # leafmap uses [latitude, longitude]
         zoom=10,
         crs='EPSG3857'
     )
 
+    # Add the mask layer to the map
     m.add_gdf(mask, layer_name="Shapefile Mask")
 
+    # Add the TIFF image to the map as RGB
     m.add_raster(tif_path, layer_name="Satellite Image", rgb=True, opacity=0.9)
 
+    # Display the map in Streamlit
     m.to_streamlit()
 
 
@@ -199,12 +188,14 @@ def main():
     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.6,
-            step=0.05
+
+        resolution = st.radio(
+
+            "Selext Processing resolution:",
+
+            (512, 1024),
+            index=0
+
         )
         overlap = st.slider(
             'Select the value of overlap',
@@ -213,8 +204,19 @@ def main():
             value=100,
             step=25
         )
-        st.write('Selected threshold value:', threshold)
+        threshold = st.slider(
+            'Select the value of the threshold',
+            min_value=0.1,
+            max_value=0.9,
+            value=0.6,
+            step=0.01
+        )
+
+        st.write('You selected:',resolution)
         st.write('Selected overlap value:', overlap)
+        st.write('Selected threshold value:', threshold)
+
+
 
         if st.button("Process File"):
             st.write("Processing...")
@@ -223,7 +225,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=resolution, overlap=overlap, batch_size=4, threshold=threshold)
 
             st.write("Processing complete!")
 
@@ -250,12 +252,14 @@ def main():
                     mime="application/zip"
                 )
 
+            # Display the map with the predicted shapefile
             display_map("output_mask.shp", "temp.tif")
 
+            # Clean up temporary files
             #os.remove("temp.tif")
             #for file in shp_files:
             #    os.remove(file)
 
 
 if __name__ == "__main__":
-    main()
+    main()