Update app.py

app.py

@@ -1,3 +1,5 @@
+
+
 import streamlit as st
 import sys
 import os
@@ -16,7 +18,7 @@ sys.path.append('Utils')
 sys.path.append('model')
 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 merge 
 from Utils.convert import read_pansharpened_rgb
@@ -34,19 +36,19 @@ if not HF_TOKEN:
 REPO_ID = "Pavan2k4/Building_area"
 REPO_TYPE = "space"
 
-# Define temporary directories outside of tracked directories
-TMP_DIR = "/tmp/space_files"
-PATCHES_DIR = os.path.join(TMP_DIR, "patches")
-PRED_PATCHES_DIR = os.path.join(TMP_DIR, "pred_patches")
-MASK_DIR = os.path.join(TMP_DIR, "generated_masks")
+# Define subdirectories using relative paths
+UPLOAD_DIR = "uploaded_images"
+MASK_DIR = "generated_masks"
+PATCHES_DIR = "patches"
+PRED_PATCHES_DIR = "pred_patches"
+CSV_LOG_PATH = "image_log.csv"
 
-# Create temporary directories
-os.makedirs(TMP_DIR, exist_ok=True)
-for directory in [PATCHES_DIR, PRED_PATCHES_DIR, MASK_DIR]:
+# Create directories
+for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
     os.makedirs(directory, exist_ok=True)
 
-# Load model and cache it to avoid unnecessary reloading
-@st.cache_resource(show_spinner=False)
+# Load model
+@st.cache_resource
 def load_model():
     model = reunet_cbam()
     model.load_state_dict(torch.load('latest.pth', map_location='cpu')['model_state_dict'])
@@ -111,8 +113,7 @@ def split(image_path, patch_size=512):
             patch_filename = f"patch_{i}_{j}.png"
             patch_path = os.path.join(PATCHES_DIR, patch_filename)
             patch.save(patch_path)
-            st.write(f"Saved patch: {patch_path}")
-
+            st.write(f"Saved patch: {patch_path}")  # Debug output
 
 def upload_page():
     if 'file_uploaded' not in st.session_state:
@@ -181,7 +182,7 @@ def upload_page():
 
                 # Count and display the number of patches
                 num_patches = len([f for f in os.listdir(PATCHES_DIR) if f.endswith('.png')])
-                #st.write(f"Number of patches created: {num_patches}")  # Debug output
+                st.write(f"Number of patches created: {num_patches}")  # Debug output
 
                 # Display buffer while analyzing
                 with st.spinner('Analyzing...'):
@@ -189,7 +190,7 @@ def upload_page():
                     for patch_filename in os.listdir(PATCHES_DIR):
                         if patch_filename.endswith(".png"):
                             patch_path = os.path.join(PATCHES_DIR, patch_filename)
-                            #st.write(f"Processing patch: {patch_path}")  # Debug output
+                            st.write(f"Processing patch: {patch_path}")  # Debug output
                             patch_img = Image.open(patch_path)
                             patch_tr_img = transforms(patch_img)
                             prediction = predict(patch_tr_img)
@@ -197,7 +198,7 @@ def upload_page():
                             mask_filename = f"mask_{patch_filename}"
                             mask_filepath = os.path.join(PRED_PATCHES_DIR, mask_filename)
                             Image.fromarray(mask).save(mask_filepath)
-                            #st.write(f"Saved mask: {mask_filepath}")  # Debug output
+                            st.write(f"Saved mask: {mask_filepath}")  # Debug output
 
                     # Merge predicted patches
                     merged_mask_filename = f"mask_{timestamp}.png"