Update app.py

app.py

@@ -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):
+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:
@@ -127,7 +127,7 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4):
                     batch_masks = model(batch_tensor.to(DEVICE))
 
                 batch_masks = torch.sigmoid(batch_masks)
-                batch_masks = (batch_masks > 0.6).float()
+                batch_masks = (batch_masks > threshold).float()
 
                 for j, mask_tensor in enumerate(batch_masks):
                     mask_resized = torch.nn.functional.interpolate(mask_tensor.unsqueeze(0),
@@ -182,6 +182,23 @@ def main():
     if uploaded_file is not None:
         st.write("File uploaded successfully!")
 
+        threshold= st.slider(
+        'Select a float value',
+        min_value=0.1,
+        max_value=0.9,
+        value=0.5,
+        step=0.05
+        )
+        overlap= st.slider(
+        'Select a float value',
+        min_value=50,
+        max_value=150,
+        value=100,
+        step=25
+        )
+        st.write('Selected threshold value:', threshold)
+        st.write('Selected overlap value:', overlap)
+
         if st.button("Process File"):
             st.write("Processing...")
 
@@ -189,7 +206,7 @@ def main():
                 f.write(uploaded_file.getbuffer())
 
             best_model.float()
-            tiles = extract_tiles("temp.tif", best_model, tile_size=512, overlap=15, batch_size=4)
+            tiles = extract_tiles("temp.tif", best_model, tile_size=512, overlap=overlap, batch_size=4,threshold=threshold)
 
             st.write("Processing complete!")