Update app.py

app.py

@@ -1,40 +1,53 @@
 import streamlit as st
-import tensorflow as tf
-from tensorflow.keras.preprocessing import image
 import numpy as np
+import cv2
+import tensorflow as tf
 from PIL import Image
+from sklearn.preprocessing import LabelEncoder  # Assuming you have this imported as well
 
 # Function to load and preprocess the uploaded image
-def load_and_preprocess_image(uploaded_image, target_size=(224, 224)):
-    img = Image.open(uploaded_image)
-    img = img.resize(target_size)
-    img_array = np.array(img)
-    img_array = np.expand_dims(img_array, axis=0)
-    img_array = img_array / 255.0
-    return img_array
-
-# Load your pre-trained model
-def load_cnn_model():
-    model = tf.keras.models.load_model('dementia_cnn_model.h5')  # Ensure correct loading
-    return model
+def load_and_preprocess_image(uploaded_file):
+    # Convert uploaded file to a format suitable for OpenCV
+    img = Image.open(uploaded_file)
+    img = img.convert("RGB")  # Convert to RGB if it's in another format
+    img = np.array(img)  # Convert to NumPy array
+    img = cv2.resize(img, (224, 224))  # Resize the image to 224x224
+    img = img / 255.0  # Normalize pixel values
+    img = np.reshape(img, (1, 224, 224, 3))  # Reshape for prediction
+    return img
+
+# Function to predict the image class
+def predict_image(img):
+    predictions = model.predict(img)  # Make a prediction
+    predicted_class_index = np.argmax(predictions[0])  # Get the predicted class index
+    return predicted_class_index
+
+# Function to get class label
+def get_class_label(predicted_class_index):
+    return label_encoder.inverse_transform([predicted_class_index])[0]  # Get class label
+
+# Load your pre-trained model (Make sure this matches the version used during training)
+model = tf.keras.models.load_model('dementia_cnn_model.h5')
 
 # Streamlit App UI
 st.title("Dementia Detection using CNN")
 st.write("Upload a brain scan (JPG format), and the model will predict its class.")
 
+# File uploader for user to upload images
 uploaded_file = st.file_uploader("Choose a JPG image...", type="jpg")
 
 if uploaded_file is not None:
     st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
-    st.write("Classifying...")
+    st.write("Detecting...")
 
+    # Load and preprocess the image
     processed_image = load_and_preprocess_image(uploaded_file)
     
-    model = load_cnn_model()  # Load the model
-    predictions = model.predict(processed_image)
-    predicted_class = np.argmax(predictions, axis=1)
+    # Make prediction
+    predicted_class_index = predict_image(processed_image)
     
-    class_names = ['Moderate Dementia', 'Very Mild Dementia', 'Mild Dementia', 'Non Demented']
-    result = class_names[predicted_class[0]]
+    # Get predicted class label
+    predicted_class_label = get_class_label(predicted_class_index)
     
-    st.write(f"Predicted Class: **{result}**")
+    # Display the result
+    st.write(f"The predicted class is: **{predicted_class_label}**")