Update app.py

app.py

@@ -1,248 +1,204 @@
-
-import streamlit as st
+import gradio as gr
 import tensorflow as tf
 from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
 import numpy as np
-import plotly.graph_objects as go
 import cv2
-from tensorflow.keras.models import Sequential
-from tensorflow.keras.layers import Dense, Dropout, Flatten
-from tensorflow.keras.optimizers import Adamax
-from tensorflow.keras.metrics import Precision, Recall
-import google.generativeai as genai
-# from google.colab import userdata
 import PIL.Image
 import os
 from dotenv import load_dotenv
-load_dotenv()
-
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-
-output_dir = 'saliency_maps'
-os.makedirs(output_dir, exist_ok=True)
-
-def generate_explanation(img_path, model_prediction, confidence):
+import traceback
+import base64
+from io import BytesIO
 
-  prompt = f"""You are an expert neurologist. You are tasked with explaining a saliency map of a brain tumor MRI scan.
-  The saliency map was generated by a deep learning model that was trained to classify brain tumors as either
-  glioma, meningioma, pituitary, or no tumor.
+# Load environment variables
+load_dotenv()
 
-  The saliency map highlights the regions of the image that the machine learning model is focusing on to make the predictions.
+# Remove Google Generative AI imports and configuration
+# We will use a local model for generating explanations
 
-  The deep learning model predicted the image to be of class '{model_prediction}' with a confidence of {confidence * 100}%.
+# Import the transformers library for text generation
+from transformers import pipeline
 
-  In your response:
-  - Explain what regions of the brain the model is focusing on, based on the saliency map. Refer to the regions highlighted in light cyan, those are the regions where the model is focusing on.
-  - Explain possible reasons why the model made the prediction it did.
-  - Don't mention anything like 'The saliency map highlights the regions the model is focusing on, which are in light cyan' in your explanation.
-  - Keep your explanation to 5 sentences max.
+# Initialize the text generation pipeline (using a smaller model for demonstration)
+# Replace 'google/flan-t5-base' with 'meta-llama/Llama-2-7b-chat-hf' if you have the resources
+generator = pipeline('text2text-generation', model='google/flan-t5-base')
 
-  Your response will go directly on the report to the doctor and patient, so don't add any extra phrases like 'Sure!' or ask any questions at the end
-  Let's think step by step about this.
-  """
+def generate_explanation(saliency_map_image, model_prediction, confidence):
+    # Convert the saliency map image array to a base64-encoded string
+    buffered = BytesIO()
+    img = PIL.Image.fromarray(saliency_map_image)
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-  img = PIL.Image.open(img_path)
+    # Prepare the prompt (Note: models like FLAN-T5 cannot process images directly)
+    prompt = f"""You are an expert neurologist. You are tasked with explaining a saliency map of a brain tumor MRI scan.
+The saliency map was generated by a deep learning model that was trained to classify brain tumors as either
+glioma, meningioma, pituitary, or no tumor.
 
-  model = genai.GenerativeModel(model_name="gemini-1.5-flash")
-  response = model.generate_content([prompt, img])
+The saliency map highlights the regions of the image that the machine learning model is focusing on to make the predictions.
 
-  return response.text
+The deep learning model predicted the image to be of class '{model_prediction}' with a confidence of {confidence * 100:.2f}%.
 
-def generate_saliency_map(model, img_array, class_index, img_size):
-  with tf.GradientTape() as tape:
-    img_tensor = tf.convert_to_tensor(img_array)
-    tape.watch(img_tensor)
-    predictions = model(img_tensor)
-    target_class = predictions[:, class_index]
 
-  gradients = tape.gradient(target_class, img_tensor)
-  gradients = tf.math.abs(gradients)
-  gradients = tf.reduce_max(gradients, axis=-1)
-  gradients = gradients.numpy().squeeze()
+In your response:
+- Explain what regions of the brain the model is focusing on, based on the saliency map.
+- Explain possible reasons why the model made the prediction it did.
+- Do not mention phrases like 'The saliency map highlights the regions the model is focusing on, which are in light cyan.'
+- Keep your explanation to 5 sentences max.
 
-  # Resize gradients to match original image size
-  gradients = cv2.resize(gradients, img_size)
+Your response will go directly in the report to the doctor and patient, so do not add extra phrases like 'Sure!' or ask any questions at the end.
+Let's think step by step about this.
+"""
 
-  # Create a circular mask for the brain area
-  center = (gradients.shape[0] // 2, gradients.shape[1] // 2)
-  radius = min(center[0], center[1]) - 10
-  y, x = np.ogrid[:gradients.shape[0], :gradients.shape[1]]
-  mask = (x - center[0])**2 + (y - center[1])**2 <= radius**2
+    # Generate the explanation using the text generation pipeline
+    response = generator(prompt, max_length=500)
+    explanation = response[0]['generated_text']
 
-  # Apply mask to gradients
-  gradients = gradients * mask
+    return explanation
 
-  # Normalize only the brain area
-  brain_gradients = gradients[mask]
-  if brain_gradients.max() > brain_gradients.min():
-    brain_gradients = (brain_gradients - brain_gradients.min()) / (brain_gradients.max() - brain_gradients.min())
-  gradients[mask] = brain_gradients
+def generate_saliency_map(model, img_array, class_index, img_size):
+    with tf.GradientTape() as tape:
+        img_tensor = tf.convert_to_tensor(img_array)
+        tape.watch(img_tensor)
+        predictions = model(img_tensor)
+        target_class = predictions[:, class_index]
 
-  # Apply a higher threshold
-  threshold = np.percentile(gradients[mask], 80)
-  gradients[gradients < threshold] = 0
+    gradients = tape.gradient(target_class, img_tensor)
+    gradients = tf.math.abs(gradients)
+    gradients = tf.reduce_max(gradients, axis=-1)
+    gradients = gradients.numpy().squeeze()
 
-  # Apply more aggressive smoothing
-  gradients = cv2.GaussianBlur(gradients, (11, 11), 0)
+    gradients = cv2.resize(gradients, img_size)
 
-  # Create a heatmap overlay with enhanced contrast
-  heatmap = cv2.applyColorMap(np.uint8(255 * gradients), cv2.COLORMAP_JET)
-  heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+    # Create a circular mask to focus on the brain region
+    center = (gradients.shape[0] // 2, gradients.shape[1] // 2)
+    radius = min(center[0], center[1]) - 10
+    y, x = np.ogrid[:gradients.shape[0], :gradients.shape[1]]
+    mask = (x - center[0])**2 + (y - center[1])**2 <= radius**2
 
-  # Resize heatmap to match original image size
-  heatmap = cv2.resize(heatmap, img_size)
+    gradients = gradients * mask
 
-  # Superimpose the heatmap on original image with increased opacity
-  original_img = image.img_to_array(img)
-  superimposed_img = heatmap * 0.7 + original_img * 0.3
-  superimposed_img = superimposed_img.astype(np.uint8)
+    # Normalize the gradients within the brain region
+    brain_gradients = gradients[mask]
+    if brain_gradients.max() > brain_gradients.min():
+        brain_gradients = (brain_gradients - brain_gradients.min()) / (brain_gradients.max() - brain_gradients.min())
+    gradients[mask] = brain_gradients
 
-  img_path = os.path.join(output_dir, uploaded_file.name)
-  with open(img_path, "wb") as f:
-    f.write(uploaded_file.getbuffer())
+    # Apply thresholding to highlight important regions
+    threshold = np.percentile(gradients[mask], 80)
+    gradients[gradients < threshold] = 0
 
-  saliency_map_path = f'saliency_maps/{uploaded_file.name}'
+    # Apply Gaussian blur to smooth the saliency map
+    gradients = cv2.GaussianBlur(gradients, (11, 11), 0)
 
-  # Save saliency map
-  cv2.imwrite(saliency_map_path, cv2.cvtColor(superimposed_img, cv2.COLOR_RGB2BGR))
+    # Create a heatmap from the gradients
+    heatmap = cv2.applyColorMap(np.uint8(255 * gradients), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
 
-  return superimposed_img
+    heatmap = cv2.resize(heatmap, img_size)
 
+    # Superimpose the heatmap on the original image
+    original_img = image.img_to_array(PIL.Image.fromarray((img_array[0] * 255).astype(np.uint8)))
+    superimposed_img = heatmap * 0.7 + original_img * 0.3
+    superimposed_img = superimposed_img.astype(np.uint8)
 
+    return superimposed_img
 
 def load_xception_model(model_path):
-  img_shape=(299, 299, 3)
-  base_model = tf.keras.applications.Xception(include_top=False, weights="imagenet", input_shape=img_shape, pooling='max')
-
-  model = Sequential([
-      base_model,
-      Flatten(),
-      Dropout(rate=0.3),
-      Dense(128, activation='relu'),
-      Dropout(rate=0.25),
-      Dense(4, activation='softmax')
-  ])
-
-  model.build((None,) + img_shape)
-
-  # Compile the model
-  model.compile(Adamax(learning_rate=0.001),
-                loss='categorical_crossentropy',
-                metrics=['accuracy', Precision(), Recall()])
-  model.load_weights(model_path)
-  return model
-
-st.title("Brain Tumor Classification")
-
-st.write("Upload an image of a brain MRI scan to classify.")
-
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-  selected_model = st.radio(
-      "Select Model",
-      ("Transfer Learning - Xception", "Custom CNN")
-  )
-
-  if selected_model == "Transfer Learning - Xception":
-    model = load_xception_model('xception_model.weights.h5')
-    img_size=(299, 299)
-  else:
-    model = load_model('cnn_model.h5')
-    img_size = (224, 224)
-
-
-  labels = ['Glioma', 'Meningioma', 'No Tumor', 'Pituitary']
-  img = image.load_img(uploaded_file, target_size=img_size)
-  img_array = image.img_to_array(img)
-  img_array = np.expand_dims(img_array, axis=0)
-  img_array /= 255.0
-
-  prediction = model.predict(img_array)
-
-  # Get the class with the highest probability
-  class_index = np.argmax(prediction[0])
-  result = labels[class_index]
-
-  saliency_map = generate_saliency_map(model, img_array, class_index, img_size)
-
-
-  col1, col2 = st.columns(2)
-  with col1:
-    st.image(uploaded_file, caption='Uploaded Image', use_column_width=True)
-  with col2:
-    st.image(saliency_map, caption='Saliency Map', use_column_width=True)
-
-
-  st.write("## Classification Results")
-
-  result_container = st.container()
-  result_container = st.container()
-  result_container.markdown(
-      f"""
-      <div style="background-color: #000000; color: #ffffff; padding: 30px; border-radius: 15px;">
-        <div style="display: flex; justify-content: space-between; align-items: center;">
-          <div style="flex: 1; text-align: center;">
-            <h3 style="color: #ffffff; margin-bottom: 10px; font-size: 20px;">Prediction</h3>
-            <p style="font-size: 36px; font-weight: 800; color: #FF0000; margin: 0;">
-              {result}
-            </p>
-          </div>
-          <div style="width: 2px; height: 80px; background-color: #ffffff; margin: 0 20px;"></div>
-          <div style="flex: 1; text-align: center;">
-            <h3 style="color: #ffffff; margin-bottom: 10px; font-size: 20px;">Confidence</h3>
-            <p style="font-size: 36px; font-weight: 800; color: #2196F3; margin: 0;">
-              {prediction[0][class_index]:.4%}
-            </p>
-          </div>
-        </div>
-      </div>
-      """,
-      unsafe_allow_html=True
-  )
-
-  # Prepare data for Plotly chart
-  probabilities = prediction[0]
-  sorted_indices = np.argsort(probabilities)[::-1]
-  sorted_labels = [labels[i] for i in sorted_indices]
-  sorted_probabilities = probabilities[sorted_indices]
-
-
-  # Create Plotly bar chart
-  fig = go.Figure(go.Bar(
-      x=sorted_probabilities,
-      y=sorted_labels,
-      orientation='h',
-      marker_color=['red' if label == result else 'blue' for label in sorted_labels]
-  ))
-
-  # Customize chart layout
-  fig.update_layout(
-      title='Probability for each class',
-      xaxis_title='Probability',
-      yaxis_title='Class',
-      height=400,
-      width=600,
-      yaxis=dict(autorange='reversed')
-  )
-
-  # Add value labels to the bars
-  for i, prob in enumerate(sorted_probabilities):
-    fig.add_annotation(
-        x=prob,
-        y=i,
-        text=f'{prob:.4f}',
-        showarrow=False,
-        xanchor='left',
-        xshift=5
+    img_shape = (299, 299, 3)
+    base_model = tf.keras.applications.Xception(include_top=False, weights="imagenet", input_shape=img_shape, pooling='max')
+
+    model = tf.keras.Sequential([
+        base_model,
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dropout(rate=0.3),
+        tf.keras.layers.Dense(128, activation='relu'),
+        tf.keras.layers.Dropout(rate=0.25),
+        tf.keras.layers.Dense(4, activation='softmax')
+    ])
+
+    model.compile(optimizer=tf.keras.optimizers.Adamax(learning_rate=0.001),
+                  loss='categorical_crossentropy',
+                  metrics=['accuracy', tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
+    model.load_weights(model_path)
+    return model
+
+def classify_brain_tumor(image_file, model_choice):
+    try:
+        # Load the selected model
+        if model_choice == "Transfer Learning - Xception":
+            model = load_xception_model('xception_model.weights.h5')
+            img_size = (299, 299)
+        else:
+            model = load_model('cnn_model.h5')
+            img_size = (224, 224)
+
+        labels = ['Glioma', 'Meningioma', 'No Tumor', 'Pituitary']
+
+        # Preprocess the input image
+        img = image.load_img(image_file, target_size=img_size)
+        img_array = image.img_to_array(img)
+        img_array = np.expand_dims(img_array, axis=0)
+        img_array /= 255.0
+
+        # Make the prediction
+        prediction = model.predict(img_array)
+        class_index = np.argmax(prediction[0])
+        result = labels[class_index]
+        confidence = prediction[0][class_index]
+
+        # Generate the saliency map
+        saliency_map = generate_saliency_map(model, img_array, class_index, img_size)
+        # No longer saving the saliency map to disk
+
+        # Generate the explanation
+        explanation = generate_explanation(saliency_map, result, confidence)
+
+        # Prepare probabilities for all classes
+        probabilities = prediction[0]
+        prob_dict = dict(zip(labels, probabilities))
+
+        # Return the outputs in the expected order
+        return [
+            result,
+            confidence,
+            saliency_map,
+            explanation,
+            "",  # Empty string for Logs
+            prob_dict  # For displaying probabilities
+        ]
+    except Exception as e:
+        # Return error information
+        return [
+            "Error",
+            0.0,
+            None,
+            "",
+            f"Error: {str(e)}\nTraceback:\n{traceback.format_exc()}",
+            {}  # Empty probabilities
+        ]
+
+def main():
+    # Define the interface
+    interface = gr.Interface(
+        fn=classify_brain_tumor,
+        inputs=[
+            gr.Image(type="filepath"),
+            gr.Radio(choices=["Transfer Learning - Xception", "Custom CNN"], label="Select Model")
+        ],
+        outputs=[
+            gr.Textbox(label="Prediction"),
+            gr.Number(label="Confidence", precision=2),
+            gr.Image(type="numpy", label="Saliency Map"),
+            gr.Textbox(label="Explanation"),
+            gr.Textbox(label="Logs"),
+            gr.Label(num_top_classes=4, label="Class Probabilities")
+        ],
+        title="Brain Tumor Classification",
+        description="Upload an MRI scan image to classify the tumor and view saliency maps with model explanations.",
     )
+    interface.launch()
 
-  # Display Plotly chart
-  st.plotly_chart(fig)
-
-  saliency_map_path = f'saliency_maps/{uploaded_file.name}'
-  explanation = generate_explanation(saliency_map_path, result, prediction[0][class_index])
-
-  st.write("## Explanation")
-  st.write(explanation)
+if __name__ == "__main__":
+    main()