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 split_merge import merge 
 from Utils.convert import read_pansharpened_rgb
 
 # Initialize Hugging Face API
@@ -34,15 +34,12 @@ if not HF_TOKEN:
 REPO_ID = "Pavan2k4/Building_area"
 REPO_TYPE = "space"
 
-# Define base directory for Hugging Face Spaces
-BASE_DIR = "/home/user"
-
-# Define subdirectories
-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")
+# 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 directories
 for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
@@ -103,6 +100,19 @@ def log_image_details(image_id, image_filename, mask_filename):
     # Save CSV to Hugging Face repo
     save_to_hf_repo(CSV_LOG_PATH, 'image_log.csv')
 
+def split(image_path, patch_size=512):
+    img = Image.open(image_path)
+    width, height = img.size
+    
+    for i in range(0, height, patch_size):
+        for j in range(0, width, patch_size):
+            box = (j, i, j+patch_size, i+patch_size)
+            patch = img.crop(box)
+            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}")  # Debug output
+
 def upload_page():
     if 'file_uploaded' not in st.session_state:
         st.session_state.file_uploaded = False
@@ -160,17 +170,25 @@ def upload_page():
             # Convert image to numpy array
             img_array = np.array(img)
 
+            st.write(f"Image shape: {img_array.shape}")  # Debug output
+
             # Check if image shape is more than 650x650
             if img_array.shape[0] > 650 or img_array.shape[1] > 650:
+                st.write("Splitting image into patches...")  # Debug output
                 # Split image into patches
                 split(converted_filepath, patch_size=512)
 
+                # 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
+
                 # Display buffer while analyzing
                 with st.spinner('Analyzing...'):
                     # Predict on each patch
                     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
                             patch_img = Image.open(patch_path)
                             patch_tr_img = transforms(patch_img)
                             prediction = predict(patch_tr_img)
@@ -178,24 +196,27 @@ 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
 
                     # Merge predicted patches
                     merged_mask_filename = f"mask_{timestamp}.png"
                     merged_mask_path = os.path.join(MASK_DIR, merged_mask_filename)
                     merge(PRED_PATCHES_DIR, merged_mask_path, img_array.shape)
+                    st.write(f"Merged mask saved: {merged_mask_path}")  # Debug output
 
                     # Save merged mask
                     st.session_state.mask_filename = merged_mask_filename
 
-                    # Clean up temporary patch files
+                    # Clean up temporary patch files but keep the folders
                     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.write("Processing whole image without splitting")  # Debug output
                 st.session_state.tr_img = transforms(img)
                 prediction = predict(st.session_state.tr_img)
                 mask = (prediction > 0.5).astype(np.uint8) * 255
@@ -203,6 +224,7 @@ def upload_page():
                 mask_filepath = os.path.join(MASK_DIR, mask_filename)
                 Image.fromarray(mask).save(mask_filepath)
                 st.session_state.mask_filename = mask_filename
+                st.write(f"Mask saved: {mask_filepath}")  # Debug output
 
             # Save mask to Hugging Face repo
             mask_filepath = os.path.join(MASK_DIR, st.session_state.mask_filename)
@@ -216,7 +238,7 @@ def upload_page():
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
             st.error("Please check the logs for more details.")
-            print(f"Error in upload_page: {str(e)}")  # This will appear in the Streamlit logs
+            st.exception(e)  # This will display the full traceback in the Streamlit app
 
     if st.session_state.file_uploaded and st.button('View result'):
         if st.session_state.filename is None:
@@ -279,8 +301,11 @@ 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)
+        # Remove files from PATCHES_DIR and PRED_PATCHES_DIR without deleting the folders
+        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.session_state.page = 'upload'
         st.session_state.file_uploaded = False
         st.session_state.filename = None