Update app.py

app.py

@@ -16,11 +16,14 @@ from shapely.ops import unary_union
 from rasterio.features import shapes
 import torch
 import numpy as np
+import tempfile
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ENCODER = 'se_resnext50_32x4d'
 ENCODER_WEIGHTS = 'imagenet'
 
+# Define a known temporary directory
+TEMP_DIR = "/tmp"
 
 #model
 @st.cache_resource
@@ -29,18 +32,14 @@ def load_model():
     model.eval().float()
     return model
 
-
 best_model = load_model()
 
-
 def to_tensor(x, **kwargs):
     return x.astype('float32')
 
-
 # Preprocessing
 preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
 
-
 def get_preprocessing(tile_size):
     _transform = [
         albu.PadIfNeeded(min_height=tile_size, min_width=tile_size, always_apply=True),
@@ -50,20 +49,13 @@ def get_preprocessing(tile_size):
     ]
     return albu.Compose(_transform)
 
-
-
-
-
 def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, threshold=0.6):
-
     preprocess = get_preprocessing(tile_size)
-
     tiles = []
 
     with rasterio.open(map_file) as src:
         height = src.height
         width = src.width
-
         effective_tile_size = tile_size - overlap
 
         for y in stqdm(range(0, height, effective_tile_size)):
@@ -122,7 +114,6 @@ def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, thres
 
     return tiles
 
-
 def create_vector_mask(tiles, output_path):
     all_polygons = []
     for mask_array, meta in tiles:
@@ -148,7 +139,6 @@ def create_vector_mask(tiles, output_path):
 
     return gdf, area_m2
 
-
 def display_map(shapefile_path, tif_path):
     st.title("Map with Shape and TIFF Overlay")
 
@@ -179,13 +169,9 @@ def display_map(shapefile_path, tif_path):
     # Display the map in Streamlit
     m.to_streamlit()
 
-
 def main():
-
-
     current_directory = os.getcwd()
     st.write('current directory:', current_directory)
-
     
     st.title("PV Segmentor")
 
@@ -194,14 +180,10 @@ def main():
     if uploaded_file is not None:
         st.write("File uploaded successfully!")
 
-
         resolution = st.radio(
-
-            "Selext Processing resolution:",
-
+            "Select Processing resolution:",
             (512, 1024),
             index=0
-
         )
         overlap = st.slider(
             'Select the value of overlap',
@@ -218,37 +200,39 @@ def main():
             step=0.01
         )
 
-        st.write('You selected:',resolution)
+        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...")
 
-            with open("temp.tif", "wb") as f:
-                f.write(uploaded_file.getbuffer())
+            # Use tempfile to create a temporary file
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.tif', dir=TEMP_DIR) as temp_file:
+                temp_filepath = temp_file.name
+                temp_file.write(uploaded_file.getbuffer())
+
+            st.write(f"Temporary file saved at: {temp_filepath}")
 
             best_model.float()
-            tiles = extract_tiles("temp.tif", best_model, tile_size=resolution, overlap=overlap, batch_size=4, threshold=threshold)
+            tiles = extract_tiles(temp_filepath, best_model, tile_size=resolution, overlap=overlap, batch_size=4, threshold=threshold)
 
             st.write("Processing complete!")
 
-            output_path = "output_mask.shp"
+            output_path = os.path.join(TEMP_DIR, "output_mask.shp")
             result_gdf, area_m2 = create_vector_mask(tiles, output_path)
 
             st.write("Vector mask created successfully!")
             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
+            shp_files = [f for f in os.listdir(TEMP_DIR) if
                          f.startswith("output_mask") and f.endswith((".shp", ".shx", ".dbf", ".prj"))]
 
             with io.BytesIO() as zip_buffer:
                 with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zip_file:
                     for file in shp_files:
-                        zip_file.write(file)
+                        zip_file.write(os.path.join(TEMP_DIR, file), file)
 
                 zip_buffer.seek(0)
                 st.download_button(
@@ -259,13 +243,14 @@ def main():
                 )
 
             # Display the map with the predicted shapefile
-            display_map("output_mask.shp", "temp.tif")
+            display_map(output_path, temp_filepath)
 
             # Clean up temporary files
-            #os.remove("temp.tif")
-            #for file in shp_files:
-            #    os.remove(file)
-
+            os.unlink(temp_filepath)
+            st.write(f"Temporary file removed: {temp_filepath}")
+            for file in shp_files:
+                os.remove(os.path.join(TEMP_DIR, file))
+            st.write("Temporary shapefile files removed")
 
 if __name__ == "__main__":
     main()