Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import tensorflow as tf
+import random
 from PIL import Image
 from tensorflow import keras
 import numpy as np
@@ -17,14 +18,19 @@ st.set_page_config(
 )
 
 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)
 
-class_names = ["Normal", "PNEUMONIA"]
+
+def prediction_cls(prediction):
+    for key, clss in class_names.items():
+        if np.argmax(prediction) == clss:
+            return key
+
 
 with st.sidebar:
     st.title("ChestAI")
@@ -42,35 +48,40 @@ with st.sidebar:
     """)
 st.set_option("deprecation.showfileUploaderEncoding", False)
 
+
 @st.cache_resource(show_spinner=False)
 def load_model():
     try:
-        from huggingface_hub import from_pretrained_keras
-        keras.utils.set_random_seed(42)
-        # Load the model from Hugging Face
-        model = from_pretrained_keras("aaravlovescodes/PneumoniaDetection")
+        from huggingface_hub import from_pretrained_keras, hf_hub_download
+        
+        # Download the model files directly
+        model_path = hf_hub_download(repo_id="ryefoxlime/PneumoniaDetection", filename="model.keras")
+        model = keras.models.load_model(model_path)
         return model
     except Exception as e:
-        st.error(f"Error loading model from Hugging Face: {str(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 from Hugging Face. Please check the model name or path.")
+    st.error("Failed to load model. Please try again.")
     st.stop()
 
 file = st.file_uploader(" ", 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 = img_array/255
 
     predictions = model.predict(img_array)
     return predictions
 
+
 if file is None:
     st.text("Please upload an image file")
 else:
@@ -78,17 +89,22 @@ else:
         image = 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",
+        ]
 
         confidence = float(max(predictions[0]) * 100)
         prediction_label = class_names[np.argmax(predictions)]
-
+        
         st.info(f"Confidence: {confidence:.2f}%")
-
+        
         if prediction_label == "Normal":
             st.balloons()
             st.success(f"Result: {prediction_label}")
         else:
             st.warning(f"Result: {prediction_label}")
-
+            
     except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
+        st.error(f"Error processing image: {str(e)}")