Upload streamlit_app.py

streamlit_app.py

@@ -1,45 +1,39 @@
 import os
+
 # Quieter TensorFlow C++ logs: 0=all, 1=warn, 2=error, 3=fatal
 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
 import streamlit as st
 from PIL import Image, ImageOps
 import tensorflow as tf
-tf.get_logger().setLevel("ERROR")
 from huggingface_hub import hf_hub_download
-
-
 from tensorflow.keras.applications.resnet50 import preprocess_input
 
+tf.get_logger().setLevel("ERROR")
 
 # ---------------- 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:
+# ---- Session state
+if "upload" not in st.session_state:
     st.session_state.upload = None
-if 'probs' 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"
 IMG_SIZE = (224, 224)
 
-
 # ---- Pretty display labels
 DISPLAY_LABELS = {
     "PO": "PO (Porosity)",
@@ -47,6 +41,8 @@ DISPLAY_LABELS = {
     "LP": "LP (Lack of Penetration)",
     "ND": "ND (No Defect)",
 }
+
+
 def pretty_label(code: str) -> str:
     return DISPLAY_LABELS.get(code, code)
 
@@ -55,32 +51,63 @@ def pretty_label(code: str) -> str:
 THRESHOLD = 0.65
 
 
+# ---------- Custom layer to handle unknown "Cast" ----------
+class CastLayer(tf.keras.layers.Layer):
+    """
+    Minimal custom layer used to replace the unknown 'Cast' layer
+    when loading the saved model from H5.
+
+    If the original object was effectively just casting to float32,
+    this reproduces that behavior.
+    """
+
+    def __init__(self, dtype="float32", **kwargs):
+        super().__init__(**kwargs)
+        self.target_dtype = tf.dtypes.as_dtype(dtype)
+
+    def call(self, inputs):
+        return tf.cast(inputs, self.target_dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"dtype": self.target_dtype.name})
+        return config
+
+
 @st.cache_resource
 def load_model_and_config_from_hub():
     """Downloads files from the Hub and loads model and config."""
-    # Get token from environment
+    # 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, 
+        repo_id=REPO_ID,
         filename=MODEL_FILENAME,
-        token=token
+        token=token,
     )
     config_path = hf_hub_download(
-        repo_id=REPO_ID, 
+        repo_id=REPO_ID,
         filename=CONFIG_FILENAME,
-        token=token
+        token=token,
     )
-    
-    # Load the Keras model
-    model = tf.keras.models.load_model(model_path, compile=False)
-    
+
+    # Load the Keras model with custom_objects so that 'Cast' is known
+    custom_objects = {
+        "Cast": CastLayer,
+    }
+
+    model = tf.keras.models.load_model(
+        model_path,
+        compile=False,
+        custom_objects=custom_objects,
+    )
+
     # Load class names from the config file
     with open(config_path, "r") as f:
         cfg = json.load(f)
-    class_names = cfg.get("class_names", ["CR", "LP", "ND", "PO"]) # Fallback
-    
+    class_names = cfg.get("class_names", ["CR", "LP", "ND", "PO"])  # Fallback
+
     return model, class_names
 
 
@@ -92,11 +119,9 @@ 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,
@@ -105,7 +130,6 @@ 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)
@@ -121,77 +145,78 @@ def upload_cb():
 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"Error loading model from Hugging Face Hub: {str(e)}")
-        st.error("Make sure the model repository is accessible.")
+        st.error("Make sure the model repository is accessible and HF_TOKEN is set if needed.")
         st.stop()
 
-
     st.file_uploader(
-        "Upload an image", type=["jpg", "jpeg", "png", "bmp", "webp"],
+        "Upload an image",
+        type=["jpg", "jpeg", "png", "bmp", "webp"],
         accept_multiple_files=False,
         on_change=upload_cb,
-        key='upload_k')
-    
+        key="upload_k",
+    )
+
     if st.session_state.upload and model is not None and class_names:
         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 = 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(
-            df[["label", "probability"]].rename(columns={"label": "Class"})
+            df[["label", "probability"]]
+            .rename(columns={"label": "Class"})
             .style.format({"probability": "{:.3f}"})
         )
 
 
 def credits():
-    st.title('Credits')
-    st.markdown('''
-[1] Benito Totino, Fanny Spagnolo, Stefania Perri, "RIAWELC: A Novel Dataset of Radiographic Images for Automatic Weld Defects Classification", in the Proceedings of the Interdisciplinary Conference on Mechanics, Computers and Electrics (ICMECE 2022), 6-7 October 2022, Barcelona, Spain.
-                
-[2] Stefania Perri, Fanny Spagnolo, Fabio Frustaci, Pasquale Corsonello, "Welding Defects Classification Through a Convolutional Neural Network", in press in Manufacturing Letters, Elsevier.
-                
+    st.title("Credits")
+    st.markdown(
+        """
+[1] Benito Totino, Fanny Spagnolo, Stefania Perri,  
+"RIAWELC: A Novel Dataset of Radiographic Images for Automatic Weld Defects Classification",  
+ICMECE 2022, Barcelona, Spain.
+
+[2] Stefania Perri, Fanny Spagnolo, Fabio Frustaci, Pasquale Corsonello,  
+"Welding Defects Classification Through a Convolutional Neural Network",  
+Manufacturing Letters, Elsevier.
+
 [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()