Update app.py

app.py

@@ -9,8 +9,7 @@ import cv2
 import numpy as np
 from PIL import Image
 import torch
-
-# Adjust import paths as needed
+from huggingface_hub import HfApi
 sys.path.append('Utils')
 sys.path.append('model')
 from model.CBAM.reunet_cbam import reunet_cbam
@@ -21,7 +20,7 @@ from split_merge import split, merge
 from Utils.convert import read_pansharpened_rgb
 
 # Define base directory for Hugging Face Spaces
-BASE_DIR = "/home/user"
+BASE_DIR = "data"
 
 # Define subdirectories
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
@@ -34,6 +33,18 @@ CSV_LOG_PATH = os.path.join(BASE_DIR, "image_log.csv")
 for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
     os.makedirs(directory, exist_ok=True)
 
+# Initialize Hugging Face API
+hf_api = HfApi()
+
+# Get the token from environment variable
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    st.error("HF_TOKEN not found in environment variables. Please set it in your Space settings.")
+    st.stop()
+
+
+REPO_ID = "Pavan2k4/Building_area" 
+
 # Load model
 @st.cache_resource
 def load_model():
@@ -49,6 +60,18 @@ 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,
+            token=HF_TOKEN
+        )
+        st.success(f"File uploaded successfully to {repo_path}")
+    except Exception as e:
+        st.error(f"Error uploading file: {str(e)}")
+
 def log_image_details(image_id, image_filename, mask_filename):
     file_exists = os.path.exists(CSV_LOG_PATH)
     
@@ -70,6 +93,9 @@ 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 upload_page():
     if 'file_uploaded' not in st.session_state:
@@ -106,6 +132,9 @@ 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...')
@@ -169,6 +198,13 @@ def upload_page():
                 Image.fromarray(mask).save(mask_filepath)
                 st.session_state.mask_filename = mask_filename
 
+            # 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:
@@ -222,7 +258,7 @@ def result_page():
         original_np = cv2.imread(original_img_path)
         mask_np = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         
-        # Ensure mask is binary
+        #  mask is binary
         _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
 
         # Resize mask to match original image size if necessary