Update app.py

app.py

@@ -27,13 +27,12 @@ st.markdown(hide_streamlit_style, unsafe_allow_html=True)
 
 
 def prediction_cls(prediction):
-    for key, clss in class_names.items():  # create a dictionary of the output classes
-        if np.argmax(prediction) == clss:  # check the class
+    for key, clss in class_names.items():
+        if np.argmax(prediction) == clss:
             return key
 
 
 with st.sidebar:
-    # st.image("mg.png")
     st.title("ChestAI")
     st.markdown("""
         ### About
@@ -50,18 +49,29 @@ with st.sidebar:
 st.set_option("deprecation.showfileUploaderEncoding", False)
 
 
-@st.cache_resource()
+@st.cache_resource(show_spinner=False)
 def load_model():
-    from huggingface_hub import from_pretrained_keras
-
-    keras.utils.set_random_seed(42)
-    model = from_pretrained_keras("ryefoxlime/PneumoniaDetection")
-    return model
-
-
-with st.spinner("Model is being loaded.."):
+    try:
+        from huggingface_hub import from_pretrained_keras
+        config = tf.compat.v1.ConfigProto()
+        config.gpu_options.allow_growth = True
+        tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))
+        
+        keras.utils.set_random_seed(42)
+        model = from_pretrained_keras("ryefoxlime/PneumoniaDetection")
+        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"])
 
 
@@ -77,22 +87,26 @@ def import_and_predict(image_data, model):
 if file is None:
     st.text("Please upload an image file")
 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)
-    np.random.seed(42)
-    x = random.randint(98, 99) + random.randint(0, 99) * 0.01
-    st.error("Accuracy : " + str(x) + " %")
-    print(predictions)
-    class_names = [
-        "Normal",
-        "PNEUMONIA",
-    ]
-
-    string = "Detected Disease : " + class_names[np.argmax(predictions)]
-    if class_names[np.argmax(predictions)] == "Normal":
-        st.balloons()
-        st.success(string)
-
-    elif class_names[np.argmax(predictions)] == "PNEUMONIA":
-        st.warning(string)
+    try:
+        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)}")