Upload streamlit_app.py

streamlit_app.py

@@ -4,6 +4,7 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 # Disable oneDNN custom ops to avoid the startup info line
 os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
+
 import json
 import numpy as np
 import pandas as pd
@@ -13,26 +14,32 @@ import tensorflow as tf
 tf.get_logger().setLevel("ERROR")
 from huggingface_hub import hf_hub_download
 
+
 from tensorflow.keras.applications.resnet50 import preprocess_input
 
+
 # ---------------- Streamlit page config ----------------
 st.set_page_config(page_title="Weld Defect Classifier", layout="centered")
 
+
 # ---- Mixed precision off on CPU to be safe
 tf.keras.mixed_precision.set_global_policy("float32")
 
 
+
 if 'upload' not in st.session_state:
     st.session_state.upload = None
 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.keras"
 CONFIG_FILENAME = "model/training_config.json"
 IMG_SIZE = (224, 224)
 
+
 # ---- Pretty display labels
 DISPLAY_LABELS = {
     "PO": "PO (Porosity)",
@@ -43,15 +50,28 @@ DISPLAY_LABELS = {
 def pretty_label(code: str) -> str:
     return DISPLAY_LABELS.get(code, code)
 
+
 # ---- Confidence threshold for displaying the prediction
 THRESHOLD = 0.65
 
+
 @st.cache_resource
 def load_model_and_config_from_hub():
     """Downloads files from the Hub and loads model and config."""
-    # Download model and config files
-    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
-    config_path = hf_hub_download(repo_id=REPO_ID, filename=CONFIG_FILENAME)
+    # Get token from environment
+    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
+    )
     
     # Load the Keras model
     model = tf.keras.models.load_model(model_path, compile=False)
@@ -63,6 +83,7 @@ def load_model_and_config_from_hub():
     
     return model, class_names
 
+
 def prepare_image(pil_img: Image.Image, target_size=(224, 224)) -> np.ndarray:
     """
     Letterbox (resize-with-pad) to target_size, fix EXIF orientation,
@@ -71,9 +92,11 @@ def prepare_image(pil_img: Image.Image, target_size=(224, 224)) -> np.ndarray:
     # 1) Honor camera EXIF orientation
     img = ImageOps.exif_transpose(pil_img)
 
+
     # 2) Convert to RGB (handles grayscale seamlessly)
     img = img.convert("RGB")
 
+
     # 3) Resize with aspect ratio preserved + pad to target (letterbox)
     img = ImageOps.pad(
         img,
@@ -82,27 +105,34 @@ def prepare_image(pil_img: Image.Image, target_size=(224, 224)) -> np.ndarray:
         color=(0, 0, 0),
     )
 
+
     # 4) To array, add batch dimension, preprocess like training
     x = np.asarray(img, dtype=np.float32)
     x = np.expand_dims(x, axis=0)
     x = preprocess_input(x)
     return x
 
+
 def upload_cb():
     st.session_state.upload = st.session_state.upload_k
     st.session_state.probs = None  # reset because the user has new input
 
+
 def weld():
     st.title("🔎 Weld Defect Classifier")
 
+
     # Load resources from the Hub
+    error_box = st.empty()
     model, class_names = None, None
     try:
         model, class_names = load_model_and_config_from_hub()
     except Exception as e:
-        st.error(f"Failed to load model: {e}")
+        st.error(f"Error loading model from Hugging Face Hub: {str(e)}")
+        st.error("Make sure the model repository is accessible.")
         st.stop()
 
+
     st.file_uploader(
         "Upload an image", type=["jpg", "jpeg", "png", "bmp", "webp"],
         accept_multiple_files=False,
@@ -113,26 +143,32 @@ def weld():
         pil_img = Image.open(st.session_state.upload)
         st.image(pil_img, caption="Input image")
 
+
         image_batch = prepare_image(pil_img, IMG_SIZE)
 
+
         if st.session_state.probs is None:
             with st.spinner("Running inference..."):
                 probs = model.predict(image_batch, verbose=0)[0].astype(float)
                 st.session_state.probs = probs
 
+
         # Build DataFrame and add pretty labels
         df = pd.DataFrame({"class": class_names, "probability": st.session_state.probs})
         df["label"] = df["class"].map(pretty_label)
         df = df.sort_values("probability", ascending=False).reset_index(drop=True)
 
+
         # Top-1 with thresholding
         top_prob = float(df.loc[0, "probability"])
         top_label = df.loc[0, "label"]
         display_label = "Unclear" if top_prob < THRESHOLD else top_label
 
+
         st.subheader("Prediction")
         st.markdown(f"**{display_label}** — Confidence: {top_prob:.3f}")
 
+
         # All probabilities
         st.subheader("All class probabilities")
         st.dataframe(
@@ -140,6 +176,7 @@ def weld():
             .style.format({"probability": "{:.3f}"})
         )
 
+
 def credits():
     st.title('Credits')
     st.markdown('''
@@ -150,9 +187,11 @@ def credits():
 [3] [Github RIAWELC](https://github.com/stefyste/RIAWELC)
     ''')
 
+
 # --- Main app navigation ---
 weld_page = st.Page(weld, title="Weld Defect Classifier", default=True)
 credit_page = st.Page(credits, title="Credits")
 
+
 pg = st.navigation([weld_page, credit_page])
 pg.run()