Update app.py

app.py

@@ -1,5 +1,3 @@
-
-
 import streamlit as st
 import sys
 import os
@@ -11,37 +9,26 @@ import cv2
 import numpy as np
 from PIL import Image
 import torch
-from huggingface_hub import HfApi
 
 # Adjust import paths as needed
 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 split_merge import split, merge 
 from Utils.convert import read_pansharpened_rgb
 
-# Initialize Hugging Face API
-hf_api = HfApi()
-
-# Get the token from secrets
-HF_TOKEN = st.secrets.get("HF_TOKEN")
-if not HF_TOKEN:
-    st.error("HF_TOKEN not found in secrets. Please set it in your Space's Configuration > Secrets.")
-    st.stop()
+# Define base directory for Hugging Face Spaces
+BASE_DIR = "/home/user"
 
-# Your Space ID (this should match exactly with your Hugging Face Space URL)
-REPO_ID = "Pavan2k4/Building_area"
-REPO_TYPE = "space"
-
-# 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"
+# 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")
 
 # Create directories
 for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
@@ -62,21 +49,6 @@ def predict(image):
         output = model(image.unsqueeze(0))
     return output.squeeze().cpu().numpy()
 
-def save_to_hf_repo(local_path, repo_path):
-    try:
-        hf_api.upload_file(
-            path_or_fileobj=local_path,
-            path_in_repo=repo_path,
-            repo_id=REPO_ID,
-            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:")
-        st.exception(e)
-
 def log_image_details(image_id, image_filename, mask_filename):
     file_exists = os.path.exists(CSV_LOG_PATH)
     
@@ -98,22 +70,6 @@ def log_image_details(image_id, image_filename, mask_filename):
             sno = 1
         
         writer.writerow([sno, date, time, 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:
@@ -150,9 +106,6 @@ def upload_page():
             
             st.success(f"Image saved to {filepath}")
             
-            # Save image to Hugging Face repo
-            save_to_hf_repo(filepath, f'uploaded_images/{filename}')
-            
             # Check if the uploaded file is a GeoTIFF
             if file_extension in ['.tiff', '.tif']:
                 st.info('Processing GeoTIFF image...')
@@ -172,25 +125,17 @@ 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)
@@ -198,27 +143,24 @@ 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 but keep the folders
+                    # Clean up temporary patch files
                     st.info('Cleaning up temporary files...')
-                    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))
+                    shutil.rmtree(PATCHES_DIR)
+                    shutil.rmtree(PRED_PATCHES_DIR)
+                    os.makedirs(PATCHES_DIR)  # Recreate empty folders
+                    os.makedirs(PRED_PATCHES_DIR)
                     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
@@ -226,21 +168,13 @@ 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)
-            save_to_hf_repo(mask_filepath, f'generated_masks/{st.session_state.mask_filename}')
-
-            # Log image details
-            log_image_details(timestamp, converted_filename, st.session_state.mask_filename)
 
             st.session_state.file_uploaded = True
             
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
             st.error("Please check the logs for more details.")
-            st.exception(e)  # This will display the full traceback in the Streamlit app
+            print(f"Error in upload_page: {str(e)}")  # This will appear in the Streamlit logs
 
     if st.session_state.file_uploaded and st.button('View result'):
         if st.session_state.filename is None:
@@ -303,11 +237,8 @@ def result_page():
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        # 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))
+        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