Update app.py

app.py

@@ -18,7 +18,6 @@ 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
@@ -34,12 +33,15 @@ if not HF_TOKEN:
 REPO_ID = "Pavan2k4/Building_area"
 REPO_TYPE = "space"
 
+# Define base directory for Hugging Face Spaces
+BASE_DIR = "DATA/"
+
 # Define subdirectories
-UPLOAD_DIR = "uploaded_images"
-MASK_DIR = "generated_masks"
-PATCHES_DIR = "patches"
-PRED_PATCHES_DIR = "pred_patches"
-CSV_LOG_PATH = "image_log.csv"
+UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
+MASK_DIR = os.path.join(BASE_DIR, "generated_masks")
+PATCHES_DIR = os.path.join(BASE_DIR, "patches")
+PRED_PATCHES_DIR = os.path.join(BASE_DIR, "pred_patches")
+CSV_LOG_PATH = os.path.join(BASE_DIR, "image_log.csv")
 
 def split(image, destination=PATCHES_DIR, patch_size=650):
     img = cv2.imread(image)
@@ -111,6 +113,7 @@ def save_to_hf_repo(local_path, repo_path):
             repo_type=REPO_TYPE,
             token=HF_TOKEN
         )
+        st.success(f"File uploaded successfully to {repo_path}")
     except Exception as e:
         st.error(f"Error uploading file: {str(e)}")
         st.error("Detailed error information:")
@@ -174,7 +177,7 @@ def upload_page():
             with open(filepath, "wb") as f:
                 f.write(bytes_data)
             
-            st.success('File uploaded and saved')
+            st.success(f"Image saved to {filepath}")
             
             # Save image to Hugging Face repo
             save_to_hf_repo(filepath, f'uploaded_images/{filename}')
@@ -184,11 +187,13 @@ def upload_page():
                 st.info('Processing GeoTIFF image...')
                 rgb_image = read_pansharpened_rgb(filepath)
                 cv2.imwrite(converted_filepath, cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR))
+                st.success(f'GeoTIFF converted to 8-bit image and saved as {converted_filename}')
                 img = Image.open(converted_filepath)
             else:
                 img = Image.open(filepath)
             
             st.image(img, caption='Uploaded Image', use_column_width=True)
+            st.success(f'Image processed and saved as {converted_filename}')
 
             # Store the full path of the converted image
             st.session_state.filename = converted_filename
@@ -225,10 +230,11 @@ def upload_page():
 
                     # 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 file in os.listdir(PATCHES_DIR):
+                        os.remove(os.path.join(PATCHES_DIR, file))
+                    for file in os.listdir(PRED_PATCHES_DIR):
+                        os.remove(os.path.join(PRED_PATCHES_DIR, file))
+                    st.success('Temporary files cleaned up')
             else:
                 # Predict on whole image
                 st.session_state.tr_img = transforms(img)
@@ -239,8 +245,6 @@ def upload_page():
                 Image.fromarray(mask).save(mask_filepath)
                 st.session_state.mask_filename = mask_filename
 
-            st.success('Mask generated and saved')
-            
             # Save mask to Hugging Face repo
             mask_filepath = os.path.join(MASK_DIR, st.session_state.mask_filename)
             save_to_hf_repo(mask_filepath, f'generated_masks/{st.session_state.mask_filename}')
@@ -255,15 +259,6 @@ 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.")
@@ -280,6 +275,8 @@ def result_page():
         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()
         return
 
@@ -323,8 +320,11 @@ def result_page():
         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.rerun()
 
 def main():
@@ -336,7 +336,4 @@ def main():
     if st.session_state.page == 'upload':
         upload_page()
     elif st.session_state.page == 'result':
-        result_page()
-
-if __name__ == '__main__':
-    main()
+        result_page()