Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,7 @@ from pathlib import Path
 from tensorflow.keras.applications import ConvNeXtLarge
 import streamlit as st
 import io
+from huggingface_hub import hf_hub_download
 
 def setup_logging():
     logging.basicConfig(
@@ -14,6 +15,32 @@ def setup_logging():
         format='%(asctime)s - %(levelname)s - %(message)s'
     )
 
+def download_model_from_hub():
+    try:
+        st.info("Model not found locally. Downloading from Hugging Face Hub...")
+        
+        # Configuration for Hugging Face Hub
+        repo_id = "Darshan03/convnext_volcano_detector"
+        filename_in_repo = "model.h5"
+        local_download_dir = "models"
+        
+        # Create models directory if it doesn't exist
+        Path(local_download_dir).mkdir(parents=True, exist_ok=True)
+        
+        # Download the model
+        local_model_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename_in_repo,
+            cache_dir=local_download_dir
+        )
+        
+        st.success(f"Model downloaded successfully to: {local_model_path}")
+        return local_model_path
+        
+    except Exception as e:
+        st.error(f"Error downloading model from Hugging Face Hub: {str(e)}")
+        raise
+
 def create_convnext_model(input_shape=(512, 512, 3)):
     base_model = ConvNeXtLarge(
         include_top=False,
@@ -120,10 +147,13 @@ def main():
     base_dir = Path(__file__).parent.absolute()
     model_path = base_dir / 'model.h5'
     
-    # Check if model exists
+    # Check if model exists, if not download from Hugging Face Hub
     if not model_path.exists():
-        st.error(f"Model not found at {model_path}")
-        return
+        try:
+            model_path = download_model_from_hub()
+        except Exception as e:
+            st.error("Failed to download model. Please check your internet connection and try again.")
+            return
     
     # Load the model
     try: