Update app.py

app.py

@@ -10,7 +10,7 @@ import numpy as np
 from PIL import Image
 import torch
 from huggingface_hub import hf_hub_download
-from datasets import Dataset, push_to_hub
+
 sys.path.append('Utils')
 sys.path.append('model')
 from model.CBAM.reunet_cbam import reunet_cbam
@@ -20,7 +20,19 @@ from Utils.area import pixel_to_sqft, process_and_overlay_image
 from Utils.convert import read_pansharpened_rgb
 
 
-dataset = Dataset("SAT_Buildings")
+from huggingface_hub import HfApi, login
+import os
+
+# Set up Hugging Face authentication
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is not set")
+
+login(token=HF_TOKEN)
+hf_api = HfApi()
+
+REPO_ID = "Pavan2k4/Building_area"
+REPO_TYPE = "space"
 
 @st.cache_resource
 def load_model():
@@ -32,8 +44,7 @@ def load_model():
 
 
 
-REPO_ID = "Pavan2k4/Building_area"
-REPO_TYPE = "space"
+
 
 # save to dir func
 def save_to_hf_repo(local_path, repo_path):
@@ -202,9 +213,14 @@ def upload_page():
 
             
             #st.success(f"Image saved to {filepath}")
-            uploaded_img = Image.open(filepath)
+            
             # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
-            dataset.add({"image":uploaded_img, "folder":"uploaded_imgs"})
+            # Save image to Hugging Face repo
+            try:
+                image_repo_path = f"images/{converted_filename}"
+                save_to_hf_repo(converted_filepath, image_repo_path)
+            except Exception as e:
+                st.error(f"Error saving image to Hugging Face repo: {str(e)}")
             
             # Check if the uploaded file is a GeoTIFF
             if file_extension in ['.tiff', '.tif']:
@@ -244,8 +260,13 @@ def upload_page():
             st.session_state.mask_filename = mask_filename
 
             # Save mask to Hugging Face repo---------------------------------------------------------------------------------------------
-            mask = Image.open(mask_filepath)
-            dataset.add({"image":mask, "folder":"generated_mask"})
+            # Save mask to Hugging Face repo
+            try:
+                mask_repo_path = f"masks/{mask_filename}"
+                save_to_hf_repo(mask_filepath, mask_repo_path)
+            except Exception as e:
+                st.error(f"Error saving mask to Hugging Face repo: {str(e)}")
+            
 
             # Log image details
             log_image_details(timestamp, converted_filename, mask_filename)