Upload streamlit_app.py

streamlit_app.py

@@ -11,7 +11,6 @@ import pandas as pd
 import streamlit as st
 from PIL import Image, ImageOps
 import tensorflow as tf
-from huggingface_hub import hf_hub_download
 from tensorflow.keras.applications.resnet50 import preprocess_input
 
 tf.get_logger().setLevel("ERROR")
@@ -28,10 +27,9 @@ if "upload" not in st.session_state:
 if "probs" not in st.session_state:
     st.session_state.probs = None
 
-# ---- Hugging Face Repo Details --- #
-REPO_ID = "zferd/welding-defect"
-MODEL_FILENAME = "model/final_single_phase.h5"
-CONFIG_FILENAME = "model/training_config.json"
+# ---- Local model file paths (inside THIS Space repo) --- #
+MODEL_PATH = "model/final_single_phase.h5"
+CONFIG_PATH = "model/training_config.json"
 IMG_SIZE = (224, 224)
 
 # ---- Pretty display labels
@@ -75,22 +73,12 @@ class CastLayer(tf.keras.layers.Layer):
 
 
 @st.cache_resource
-def load_model_and_config_from_hub():
-    """Downloads files from the Hub and loads model and config."""
-    # Get token from environment (set in HF Space secrets if repo is private)
-    token = os.environ.get("HF_TOKEN")
-
-    # Download model and config files with token
-    model_path = hf_hub_download(
-        repo_id=REPO_ID,
-        filename=MODEL_FILENAME,
-        token=token,
-    )
-    config_path = hf_hub_download(
-        repo_id=REPO_ID,
-        filename=CONFIG_FILENAME,
-        token=token,
-    )
+def load_model_and_config():
+    """Loads model + config from local files inside the Space."""
+    if not os.path.exists(MODEL_PATH):
+        raise FileNotFoundError(f"Model file not found at: {MODEL_PATH}")
+    if not os.path.exists(CONFIG_PATH):
+        raise FileNotFoundError(f"Config file not found at: {CONFIG_PATH}")
 
     # Load the Keras model with custom_objects so that 'Cast' is known
     custom_objects = {
@@ -98,13 +86,13 @@ def load_model_and_config_from_hub():
     }
 
     model = tf.keras.models.load_model(
-        model_path,
+        MODEL_PATH,
         compile=False,
         custom_objects=custom_objects,
     )
 
     # Load class names from the config file
-    with open(config_path, "r") as f:
+    with open(CONFIG_PATH, "r") as f:
         cfg = json.load(f)
     class_names = cfg.get("class_names", ["CR", "LP", "ND", "PO"])  # Fallback
 
@@ -145,14 +133,13 @@ def upload_cb():
 def weld():
     st.title("🔎 Weld Defect Classifier")
 
-    # Load resources from the Hub
-    model, class_names = None, None
+    # Load resources from local files
     try:
-        model, class_names = load_model_and_config_from_hub()
+        model, class_names = load_model_and_config()
     except Exception as e:
-        st.error(f"Error loading model from Hugging Face Hub: {str(e)}")
-        st.error("Make sure the model repository is accessible and HF_TOKEN is set if needed.")
+        st.error(f"Error loading model/config: {str(e)}")
         st.stop()
+        return
 
     st.file_uploader(
         "Upload an image",