Update app.py

app.py

@@ -226,12 +226,17 @@ 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...')
-                    shutil.rmtree(PATCHES_DIR)
-                    shutil.rmtree(PRED_PATCHES_DIR)
-                    os.makedirs(PATCHES_DIR)  # Recreate empty folders
-                    os.makedirs(PRED_PATCHES_DIR)
+                    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)
             else:
                 # Predict on whole image
                 st.session_state.tr_img = transforms(img)
@@ -319,12 +324,7 @@ def result_page():
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        shutil.rmtree(PATCHES_DIR)
-        shutil.rmtree(PRED_PATCHES_DIR)
-        st.session_state.page = 'upload'
-        st.session_state.file_uploaded = False
-        st.session_state.filename = None
-        st.session_state.mask_filename = None
+        
         st.rerun()
 
 def main():