Update app.py

app.py

@@ -10,10 +10,7 @@ import numpy as np
 from PIL import Image
 import torch
 from huggingface_hub import hf_hub_download
-from huggingface_hub import HfApi
-
-
-
+from dataset import Dataset, push_to_hub
 sys.path.append('Utils')
 sys.path.append('model')
 from model.CBAM.reunet_cbam import reunet_cbam
@@ -22,6 +19,9 @@ from model.unet import UNET
 from Utils.area import pixel_to_sqft, process_and_overlay_image
 from Utils.convert import read_pansharpened_rgb
 
+
+dataset = Dataset("SAT_Buildings")
+
 @st.cache_resource
 def load_model():
     model = reunet_cbam()
@@ -29,16 +29,6 @@ def load_model():
     model.eval()
     return model
 
-@st.cache_resource
-def get_api():
-    hf_api = HfApi()
-    HF_TOKEN = st.secrets.get("HF_TOKEN")
-    if not HF_TOKEN:
-        st.error("HF_TOKEN not found")
-        st.stop()
-    return HF_TOKEN,hf_api
-
-HF_TOKEN,hf_api = get_api()
 
 
 
@@ -212,9 +202,9 @@ def upload_page():
 
             
             #st.success(f"Image saved to {filepath}")
-            
+            uploaded_img = Image.open(filepath)
             # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
-            save_to_hf_repo(filepath, f'uploaded_images/{filename}')
+            dataset.add({"image":uploaded_img, "folder":"uploaded_imgs"})
             
             # Check if the uploaded file is a GeoTIFF
             if file_extension in ['.tiff', '.tif']:
@@ -254,7 +244,8 @@ def upload_page():
             st.session_state.mask_filename = mask_filename
 
             # Save mask to Hugging Face repo---------------------------------------------------------------------------------------------
-            save_to_hf_repo(mask_filepath, f'generated_masks/{st.session_state.mask_filename}')
+            mask = Image.open(mask_filepath)
+            dataset.add({"image":mask, "folder":"generated_mask"})
 
             # Log image details
             log_image_details(timestamp, converted_filename, mask_filename)