Update app.py

app.py

@@ -18,7 +18,7 @@ from model.CBAM.reunet_cbam import reunet_cbam
 from model.transform import transforms
 from model.unet import UNET
 from Utils.area import pixel_to_sqft, process_and_overlay_image
-
+from split_merge import split, merge 
 from Utils.convert import read_pansharpened_rgb
 
 # Initialize Hugging Face API
@@ -40,9 +40,7 @@ MASK_DIR = "generated_masks"
 PATCHES_DIR = "patches"
 PRED_PATCHES_DIR = "pred_patches"
 CSV_LOG_PATH = "image_log.csv"
-# Create directories
-for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
-    os.makedirs(directory, exist_ok=True)
+
 
 def split(image, destination = PATCHES_DIR, patch_size = 650):
     img = cv2.imread(image)
@@ -88,7 +86,9 @@ def merge(patch_folder , dest_image = 'out.png', image_shape = None):
 
     cv2.imwrite(dest_image, merged)
 
-
+# Create directories
+for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
+    os.makedirs(directory, exist_ok=True)
 
 # Load model
 @st.cache_resource
@@ -156,12 +156,6 @@ def upload_page():
 
     image = st.file_uploader('Choose a satellite image', type=['jpg', 'png', 'jpeg', 'tiff', 'tif'])
 
-    if st.button('Back to Upload'):
-        st.session_state.page = 'upload'
-        st.session_state.file_uploaded = False
-        st.session_state.filename = None
-        st.session_state.mask_filename = None
-        st.rerun()
     if image is not None and not st.session_state.file_uploaded:
         try:
             bytes_data = image.getvalue()
@@ -232,17 +226,12 @@ def upload_page():
                     # Save merged mask
                     st.session_state.mask_filename = merged_mask_filename
 
-                    # Clean up temporary patch files
                     # Clean up temporary patch files
                     st.info('Cleaning up temporary files...')
-                    for filename in os.listdir(PATCHES_DIR):
-                        file_path = os.path.join(PATCHES_DIR, filename)
-                        if os.path.isfile(file_path):
-                            os.remove(file_path)
-                    for filename in os.listdir(PRED_PATCHES_DIR):
-                        file_path = os.path.join(PRED_PATCHES_DIR, filename)
-                        if os.path.isfile(file_path):
-                            os.remove(file_path)
+                    shutil.rmtree(PATCHES_DIR)
+                    shutil.rmtree(PRED_PATCHES_DIR)
+                    os.makedirs(PATCHES_DIR)  # Recreate empty folders
+                    os.makedirs(PRED_PATCHES_DIR)
             else:
                 # Predict on whole image
                 st.session_state.tr_img = transforms(img)
@@ -269,6 +258,15 @@ def upload_page():
             st.error("Please check the logs for more details.")
             print(f"Error in upload_page: {str(e)}")  # This will appear in the Streamlit logs
 
+    for filename in os.listdir(PATCHES_DIR):
+    file_path = os.path.join(PATCHES_DIR, filename)
+        if os.path.isfile(file_path):
+            os.remove(file_path)
+    for filename in os.listdir(PRED_PATCHES_DIR):
+    file_path = os.path.join(PRED_PATCHES_DIR, filename)
+        if os.path.isfile(file_path):
+            os.remove(file_path)
+
     if st.session_state.file_uploaded and st.button('View result'):
         if st.session_state.filename is None:
             st.error("Please upload an image before viewing the result.")
@@ -287,7 +285,7 @@ def result_page():
             st.session_state.file_uploaded = False
             st.session_state.filename = None
             st.session_state.mask_filename = None
-            st.experimental_rerun()
+            st.rerun()
         return
 
     col1, col2 = st.columns(2)
@@ -329,13 +327,13 @@ def result_page():
     else:
         st.error("Image or mask file not found for overlay.")
 
-    
     if st.button('Back to Upload'):
+        
         st.session_state.page = 'upload'
         st.session_state.file_uploaded = False
         st.session_state.filename = None
         st.session_state.mask_filename = None
-        st.experimental_rerun()
+        st.rerun()
 
 def main():
     st.title('Building area estimation')