Update app.py

app.py

@@ -1,65 +1,142 @@
-import os
 import tensorflow as tf
 import numpy as np
 import streamlit as st
+import matplotlib.pyplot as plt
 from PIL import Image
 from fpdf import FPDF
 
-# Define the model path
-MODEL_URL = "https://huggingface.co/chrisaldikaraharja/BrainTumor-Model/resolve/main/brain_tumor_detection_model.h5"
-MODEL_PATH = "brain_tumor_detection_model.h5"
+# Load the pre-trained model from Hugging Face
+model_url = "https://huggingface.co/chrisaldikaraharja/BrainTumor-Model/resolve/main/brain_tumor_detection_model.h5"
 
-# Function to download the model if it doesn't exist
-def download_model(model_path, model_url):
-    if not os.path.exists(model_path):
-        tf.keras.utils.get_file(model_path, model_url)
+# Function to load the model
+def load_model(model_path):
+    model = tf.keras.models.load_model(model_path, compile=False)
+    return model
 
-# Download the model
-download_model(MODEL_PATH, MODEL_URL)
-
-# Load the model
+# Attempt to load the model
 try:
-    model = tf.keras.models.load_model(MODEL_PATH, compile=False)
+    model = load_model(model_url)
+    st.success("Model loaded successfully.")
 except Exception as e:
     st.error(f"Error encountered while loading the model: {e}")
     model = None
 
-# Check if model is loaded successfully
-if model is None:
-    st.stop()
-
 # Define the class labels
 label_map = {0: 'Glioma', 1: 'Meningioma', 2: 'Normal', 3: 'Pituitary'}
 
 def preprocess_image(image):
-    img = image.resize((150, 150))
-    img = np.array(img) / 255.0
-    img = np.expand_dims(img, axis=0)
-    return img
+    try:
+        img = image.resize((150, 150))  # Resize to match model input shape
+        img = np.array(img) / 255.0      # Normalize pixel values
+        img = np.expand_dims(img, axis=0)  # Add batch dimension
+        return img
+    except Exception as e:
+        st.error(f"Error in preprocessing image: {e}")
+        raise
+
+# Function to compute Occlusion Sensitivity
+def compute_occlusion_sensitivity(model, img_array, class_idx, patch_size=15):
+    img_array_np = img_array.squeeze()
+    img_shape = img_array_np.shape
+
+    original_img = np.copy(img_array_np)
+    original_prediction = model.predict(np.expand_dims(original_img, axis=0))[0, class_idx]
+
+    sensitivity_map = np.zeros((img_shape[0], img_shape[1]))
+
+    for i in range(0, img_shape[0], patch_size):
+        for j in range(0, img_shape[1], patch_size):
+            img_copy = np.copy(original_img)
+            img_copy[i:i+patch_size, j:j+patch_size, :] = 0  # Occlude patch
+
+            occluded_prediction = model.predict(np.expand_dims(img_copy, axis=0))[0, class_idx]
+            sensitivity_map[i:i+patch_size, j:j+patch_size] = original_prediction - occluded_prediction
+
+    return sensitivity_map
 
-# Main prediction function
-def predict(image):
+# Function to visualize the occlusion map
+def visualize_occlusion_map(img, occlusion_map):
+    plt.figure(figsize=(10, 10))
+    plt.imshow(img)
+    plt.imshow(occlusion_map, cmap='hot', alpha=0.5)
+    plt.axis('off')
+    plt.colorbar()
+    occlusion_path = "/tmp/occlusion_map.png"
+    plt.savefig(occlusion_path)
+    plt.close()
+    return occlusion_path
+
+# Function to generate PDF report
+def create_pdf_report(prediction_text, doctor_notes, image_path):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, prediction_text)
+    pdf.ln(10)
+    pdf.set_font("Arial", 'B', size=12)
+    pdf.cell(0, 10, "Doctor's Notes:")
+    pdf.ln(10)
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, doctor_notes)
+    pdf.image(image_path, x=10, y=pdf.get_y(), w=pdf.w - 20)
+    pdf_output_path = "/tmp/medical_report.pdf"
+    pdf.output(pdf_output_path)
+    return pdf_output_path
+
+# Main prediction and report generation function
+def predict_and_generate_report(image, doctor_notes):
     try:
         img_array = preprocess_image(image)
+        print(f"Preprocessed image shape: {img_array.shape}")  # Debugging output
+        
         predictions = model.predict(img_array)[0]
         class_idx = np.argmax(predictions)
         class_label = label_map[class_idx]
 
+        # Prediction text
         prediction_results = [f"{label_map[i]}: {predictions[i] * 100:.2f}%" for i in range(len(predictions))]
         prediction_text = f"Prediction: {class_label}\n" + "\n".join(prediction_results)
+        print(f"Predictions: {predictions}")
+
+        # Save the uploaded image
+        image_path = "/tmp/image.png"
+        image.save(image_path)
+
+        # Generate PDF report
+        pdf_report_path = create_pdf_report(prediction_text, doctor_notes, image_path)
+
+        # Compute and visualize occlusion sensitivity
+        occlusion_map = compute_occlusion_sensitivity(model, img_array, class_idx)
+        occlusion_map_path = visualize_occlusion_map(np.array(image), occlusion_map)
+
+        return image, prediction_text, pdf_report_path, occlusion_map_path
 
-        return prediction_text
     except Exception as e:
-        return f"Error during prediction: {e}"
+        print(f"Error during prediction: {e}")
+        return None, "Error during prediction. Please check the input image.", None, None
+
+# Streamlit Interface
+st.title("Brain Tumor Detection and Occlusion Sensitivity")
+st.write("Upload an MRI image to detect brain tumors and view which areas contributed to the prediction.")
+
+# Image upload
+uploaded_image = st.file_uploader("Upload MRI Image", type=["png", "jpg", "jpeg"])
+doctor_notes = st.text_area("Enter doctor's notes here...", "")
+
+if st.button("Predict"):
+    if uploaded_image is not None:
+        image = Image.open(uploaded_image)
+        drawn_image, prediction_text, pdf_report_path, occlusion_map_path = predict_and_generate_report(image, doctor_notes)
 
-# Streamlit UI
-st.title("Brain Tumor Detection")
-uploaded_file = st.file_uploader("Upload MRI Image", type=["png", "jpg", "jpeg"])
+        if drawn_image is not None:
+            st.image(drawn_image, caption="Original MRI Image", use_column_width=True)
+            st.text(prediction_text)
 
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
+            # Provide PDF report download
+            st.download_button("Download Medical Report", data=open(pdf_report_path, "rb").read(), file_name="medical_report.pdf")
 
-    if st.button("Predict"):
-        prediction_text = predict(image)
-        st.write(prediction_text)
+            # Display Occlusion Sensitivity Map
+            occlusion_image = Image.open(occlusion_map_path)
+            st.image(occlusion_image, caption="Occlusion Sensitivity Map", use_column_width=True)
+    else:
+        st.warning("Please upload an MRI image to proceed.")