fix: Update model loading logic for Hugging Face Spaces

frontend/app.py

@@ -56,16 +56,19 @@ st.title("Medi Scape Dashboard")
 # --- Session State Initialization ---
 if 'disease_model' not in st.session_state:
     try:
-        model_path = 'FINAL_MODEL.zip'  # Updated path to zip file
-        print(f"Attempting to load disease model from: {model_path}")
-        print(f"Model file exists: {os.path.exists(model_path)}") 
-        with tf.keras.utils.get_file('FINAL_MODEL.keras', model_path, extract=True) as extracted_model_path:
-            model_dir = os.path.dirname(extracted_model_path)
-            model_path = os.path.join(model_dir, 'FINAL_MODEL.keras')
+        # Check if running on Hugging Face Spaces
+        if 'HF_SPACE_ID' in os.environ:
+            model_path = 'FINAL_MODEL.zip'
+            with tf.keras.utils.get_file('FINAL_MODEL.keras', model_path, extract=True) as extracted_model_path:
+                model_dir = os.path.dirname(extracted_model_path)
+                model_path = os.path.join(model_dir, 'FINAL_MODEL.keras')
+                st.session_state.disease_model = tf.keras.models.load_model(model_path)
+        else:  # Assume running locally
+            model_path = 'FINAL_MODEL.keras'
             st.session_state.disease_model = tf.keras.models.load_model(model_path)
         print("Disease model loaded successfully!")
     except FileNotFoundError:
-        st.error("Disease classification model not found. Please ensure 'FINAL_MODEL.zip' is in the same directory as this app.")
+        st.error("Disease classification model not found. Please ensure 'FINAL_MODEL.zip' (for Hugging Face) or 'FINAL_MODEL.keras' (for local) is in the correct directory.")
         st.session_state.disease_model = None
     except PermissionError:
         st.error("Permission error accessing 'model.weights.h5'. Please ensure the file is not being used by another process.")
@@ -102,15 +105,20 @@ if 'model_llm' not in st.session_state:
 
 # --- End of Session State Initialization ---
 
-# Load the disease classification model
+# Load the disease classification model (outside session state for this example)
 try:
-    model_path = 'FINAL_MODEL.zip'  # Updated path to zip file
-    with tf.keras.utils.get_file('FINAL_MODEL.keras', model_path, extract=True) as extracted_model_path:
-        model_dir = os.path.dirname(extracted_model_path)
-        model_path = os.path.join(model_dir, 'FINAL_MODEL.keras')
+    # Check if running on Hugging Face Spaces
+    if 'HF_SPACE_ID' in os.environ:
+        model_path = 'FINAL_MODEL.zip'
+        with tf.keras.utils.get_file('FINAL_MODEL.keras', model_path, extract=True) as extracted_model_path:
+            model_dir = os.path.dirname(extracted_model_path)
+            model_path = os.path.join(model_dir, 'FINAL_MODEL.keras')
+            disease_model = tf.keras.models.load_model(model_path)
+    else:  # Assume running locally
+        model_path = 'FINAL_MODEL.keras'
         disease_model = tf.keras.models.load_model(model_path)
 except FileNotFoundError:
-    st.error("Disease classification model not found. Please ensure 'FINAL_MODEL.zip' is in the same directory as this app.")
+    st.error("Disease classification model not found. Please ensure 'FINAL_MODEL.zip' (for Hugging Face) or 'FINAL_MODEL.keras' (for local) is in the correct directory.")
     disease_model = None
 except PermissionError:
     st.error("Permission error accessing 'model.weights.h5'. Please ensure the file is not being used by another process.")
@@ -401,7 +409,7 @@ else:
         st.write("Upload a chest X-ray image for disease detection.")
         uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-        if uploaded_image is not None and st.session_state.disease_model is not None:
+        if uploaded_image is not None and disease_model is not None:  # Use disease_model directly
             # Display the image
             img_opened = Image.open(uploaded_image).convert('RGB')
             image_pred = np.array(img_opened)
@@ -417,7 +425,7 @@ else:
             image_pred = np.expand_dims(image_pred, axis=0)
 
             # Predict using the model
-            prediction = st.session_state.disease_model.predict(image_pred)
+            prediction = disease_model.predict(image_pred)  # Use disease_model directly
 
             # Get the predicted class
             predicted_ = np.argmax(prediction)