Update app.py

app.py

@@ -16,30 +16,31 @@ 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
+# Load and prepare the model
 @st.cache_resource
 def load_model():
     model = torch.load('deeplabv3 v15.pth', map_location=DEVICE)
     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),
@@ -49,13 +50,20 @@ 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)):
@@ -114,6 +122,7 @@ 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:
@@ -139,6 +148,7 @@ 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")
 
@@ -169,10 +179,8 @@ 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")
 
     uploaded_file = st.file_uploader("Choose a TIF file", type="tif")
@@ -180,10 +188,14 @@ def main():
     if uploaded_file is not None:
         st.write("File uploaded successfully!")
 
+
         resolution = st.radio(
-            "Select Processing resolution:",
+
+            "Selext Processing resolution:",
+
             (512, 1024),
             index=0
+
         )
         overlap = st.slider(
             'Select the value of overlap',
@@ -200,39 +212,37 @@ 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...")
 
-            # 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}")
+            with open("temp.tif", "wb") as f:
+                f.write(uploaded_file.getbuffer())
 
             best_model.float()
-            tiles = extract_tiles(temp_filepath, best_model, tile_size=resolution, 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!")
 
-            output_path = os.path.join(TEMP_DIR, "output_mask.shp")
+            output_path = "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(TEMP_DIR) if
+            shp_files = [f for f in os.listdir() 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(os.path.join(TEMP_DIR, file), file)
+                        zip_file.write(file)
 
                 zip_buffer.seek(0)
                 st.download_button(
@@ -243,14 +253,13 @@ def main():
                 )
 
             # Display the map with the predicted shapefile
-            display_map(output_path, temp_filepath)
+            display_map("output_mask.shp", "temp.tif")
 
             # Clean up temporary files
-            #os.unlink(temp_filepath)
-            #st.write(f"Temporary file removed: {temp_filepath}")
+            #os.remove("temp.tif")
             #for file in shp_files:
-                #os.remove(os.path.join(TEMP_DIR, file))
-            #st.write("Temporary shapefile files removed")
+            #    os.remove(file)
+
 
 if __name__ == "__main__":
     main()