Update app.py

app.py

@@ -1,80 +1,50 @@
 import streamlit as st
-import numpy as np
 import tensorflow as tf
+import numpy as np
 from PIL import Image
-import nibabel as nib
-from scipy import ndimage
+from tensorflow.keras.applications.densenet import preprocess_input
+
+# Define class labels
+CLASS_LABELS = [
+    "Bacterial Pneumonia",
+    "Corona Virus Disease",
+    "Normal",
+    "Tuberculosis",
+    "Viral Pneumonia"
+]
 
 # Load the trained model
-@st.cache_resource()
+@st.cache_resource
 def load_model():
-    return tf.keras.models.load_model("model.h5")
+    model = tf.keras.models.load_model("lung_disease_model.h5")  # Update filename if needed
+    return model
 
 model = load_model()
 
-# Preprocessing functions
-def read_nifti_file(filepath):
-    """Read and load volume"""
-    scan = nib.load(filepath)
-    scan = scan.get_fdata()
-    return scan
-
-def normalize(volume):
-    """Normalize the volume"""
-    volume = volume.astype(np.float32)  # Ensure float32 before operations
-    min_val = -1000
-    max_val = 400
-    volume[volume < min_val] = min_val
-    volume[volume > max_val] = max_val
-    volume = (volume - min_val) / (max_val - min_val)
-    return volume
-
-
-def resize_volume(img):
-    """Resize across z-axis"""
-    desired_depth = 64
-    desired_width = 128
-    desired_height = 128
-    current_depth = img.shape[-1]
-    current_width = img.shape[0]
-    current_height = img.shape[1]
-    depth = current_depth / desired_depth
-    width = current_width / desired_width
-    height = current_height / desired_height
-    depth_factor = 1 / depth
-    width_factor = 1 / width
-    height_factor = 1 / height
-    img = ndimage.rotate(img, 90, reshape=False)
-    img = ndimage.zoom(img, (width_factor, height_factor, depth_factor), order=1)
-    return img
-
-
-def process_scan(image):
-    """Process the image"""
-    volume = np.array(image)
-    volume = normalize(volume)
-    volume = resize_volume(volume)
-    volume = np.expand_dims(volume, axis=-1)  # Add channel dimension
-    return volume
+# Function to preprocess image
+def preprocess_image(image):
+    image = image.resize((224, 224))  # Match your model's input size
+    image = np.array(image)  # Convert to NumPy array
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    image = preprocess_input(image)  # Apply DenseNet preprocessing
+    return image
 
-# Define the Streamlit app
-st.title("Pneumonia CT Scan")
-st.write("Upload an image to classify it using the trained model.")
+# Streamlit UI
+st.title("Lung Disease Classification")
 
-uploaded_file = st.file_uploader("Choose an image...")
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
 
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
     
-    # Preprocess the image
-    img_array = process_scan(image)  # Apply preprocessing pipeline
-    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    # Preprocess and predict
+    processed_image = preprocess_image(image)
+    prediction = model.predict(processed_image)
     
-    # Make a prediction
-    prediction = model.predict(img_array)
-    predicted_class = np.argmax(prediction, axis=1)[0]
-    confidence = np.max(prediction)
+    # Get the predicted class index
+    predicted_class = np.argmax(prediction)
+    predicted_label = CLASS_LABELS[predicted_class]
     
-    st.write(f"### Prediction: Class {predicted_class}")
-    st.write(f"### Confidence: {confidence:.2f}")
+    # Display result
+    st.write(f"### Prediction: **{predicted_label}**")