Update app.py

app.py

@@ -1,37 +1,53 @@
 import streamlit as st
 import tensorflow as tf
-import random
-from PIL import Image
-from tensorflow import keras
 import numpy as np
+from PIL import Image
+from huggingface_hub import hf_hub_download
 import os
 
+# Suppress warnings
 import warnings
-
 warnings.filterwarnings("ignore")
-os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 
+# Set page configuration
 st.set_page_config(
     page_title="ChestAI - Pneumonia Detection",
     page_icon="🫁",
     initial_sidebar_state="auto",
 )
 
+# Hide Streamlit style
 hide_streamlit_style = """
-	<style>
-  #MainMenu {visibility: hidden;}
-	footer {visibility: hidden;}
-  </style>
+<style>
+#MainMenu {visibility: hidden;}
+footer {visibility: hidden;}
+</style>
 """
 st.markdown(hide_streamlit_style, unsafe_allow_html=True)
 
+# Function to load the model
+@st.cache_resource(show_spinner=False)
+def load_model():
+    try:
+        # Download the model directory
+        model_dir = hf_hub_download(repo_id="ryefoxlime/PneumoniaDetection", repo_type="model", library="tf", cache_dir="/home/user/.cache/huggingface/hub")
 
-def prediction_cls(prediction):
-    for key, clss in class_names.items():
-        if np.argmax(prediction) == clss:
-            return key
+        # Load the model using tf.saved_model.load
+        model = tf.saved_model.load(model_dir)
+        return model
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None
 
+# Load the model
+with st.spinner("Model is being loaded..."):
+    model = load_model()
+
+if model is None:
+    st.error("Failed to load model. Please try again.")
+    st.stop()
 
+# Sidebar for app information
 with st.sidebar:
     st.title("ChestAI")
     st.markdown("""
@@ -46,68 +62,43 @@ with st.sidebar:
         ### Note
         This tool is for educational purposes only. Always consult healthcare professionals for medical advice.
     """)
-st.set_option("deprecation.showfileUploaderEncoding", False)
-
 
-@st.cache_resource(show_spinner=False)
-def load_model():
-    try:
-        from huggingface_hub import hf_hub_download
-        from keras.layers import TFSMLayer
-        
-        # Download the model files directly
-        model_path = hf_hub_download(repo_id="ryefoxlime/PneumoniaDetection", filename="saved_model.pb")
-        
-        # Use TFSMLayer to load the model
-        model = TFSMLayer(model_path, call_endpoint='serving_default')
-        return model
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None
-
-
-with st.spinner("Model is being loaded..."):
-    model = load_model()
-
-if model is None:
-    st.error("Failed to load model. Please try again.")
-    st.stop()
-
-file = st.file_uploader(" ", type=["jpg", "png"])
+st.set_option("deprecation.showfileUploaderEncoding", False)
 
+# File uploader for image input
+file = st.file_uploader("Upload a chest X-ray image", type=["jpg", "png"])
 
 def import_and_predict(image_data, model):
-    img_array = keras.preprocessing.image.img_to_array(image_data)
-    img_array = np.expand_dims(img_array, axis=0)
-    img_array = img_array/255
+    img_array = tf.keras.preprocessing.image.img_to_array(image_data)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    img_array = img_array / 255.0  # Normalize the image
 
-    predictions = model.predict(img_array)
+    # Perform prediction
+    predictions = model(img_array)  # Call the model for prediction
     return predictions
 
+# Class names for prediction results
+class_names = ["Normal", "PNEUMONIA"]
 
 if file is None:
     st.text("Please upload an image file")
 else:
     try:
-        image = keras.preprocessing.image.load_img(file, target_size=(224, 224), color_mode='rgb')
+        image = tf.keras.preprocessing.image.load_img(file, target_size=(224, 224), color_mode='rgb')
         st.image(image, caption="Uploaded Image.", use_column_width=True)
-        predictions = import_and_predict(image, model)
         
-        class_names = [
-            "Normal",
-            "PNEUMONIA",
-        ]
+        predictions = import_and_predict(image, model)
+        predicted_class = np.argmax(predictions)  # Get the index of the highest prediction
+        confidence = float(predictions[0][predicted_class] * 100)  # Confidence percentage
 
-        confidence = float(max(predictions[0]) * 100)
-        prediction_label = class_names[np.argmax(predictions)]
-        
+        # Display the results
         st.info(f"Confidence: {confidence:.2f}%")
         
-        if prediction_label == "Normal":
+        if class_names[predicted_class] == "Normal":
             st.balloons()
-            st.success(f"Result: {prediction_label}")
+            st.success(f"Result: {class_names[predicted_class]}")
         else:
-            st.warning(f"Result: {prediction_label}")
-            
+            st.warning(f"Result: {class_names[predicted_class]}")
+
     except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
+        st.error(f"Error processing image: {str(e)}")