Update app.py

app.py

@@ -1,97 +1,51 @@
+#!/usr/bin/env python3
+"""
+Streamlit Brain MRI Tumor Detection App (updated, safe startup + LLM safety)
+- Monkeypatches torch.classes.__path__ before importing streamlit to avoid a Streamlit <-> PyTorch watcher crash.
+- Returns probabilistic model output (label + confidence).
+- Adds a visible medical disclaimer.
+- Adds robust error handling around the Groq (Deepseek R1) call and ensures the LLM output contains a safety sentence.
+- Keeps your original UI/CSS with small UX improvements.
+"""
+
+import os
+import logging
+import traceback
+
+# ------------------ Safe startup: import torch first and monkeypatch ------------------
+# This avoids Streamlit's file-watcher triggering PyTorch C++ registry introspection errors.
+try:
+    import torch
+    # Force a benign __path__ so Streamlit's watcher won't attempt unsafe introspection.
+    try:
+        # If torch.classes exists, ensure __path__ is present and is a harmless list.
+        if hasattr(torch, "classes"):
+            # Some torch builds may already have __path__; overwrite safely.
+            torch.classes.__path__ = []
+    except Exception as _e:
+        # If something goes wrong, don't crash the app at module import time.
+        logging.warning("Failed to set torch.classes.__path__: %s", _e)
+except Exception as e:
+    # If torch can't be imported at all, we still continue so Streamlit can display an error to the user.
+    # Log it; later we'll surface a friendly message in the UI.
+    logging.error("Unable to import torch at startup: %s\n%s", e, traceback.format_exc())
+    torch = None
+
+# ------------------ Now safe to import Streamlit and other packages ------------------
 import streamlit as st
 from PIL import Image
-import torch
-from transformers import ViTForImageClassification, ViTImageProcessor
-import logging
-import base64
 from io import BytesIO
-from groq import Groq  # Import the Groq client for Deepseek R1 API
-import os
-# ------------------ Setup Logging ------------------
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# ------------------ Load the ViT Model ------------------
-repository_id = "EnDevSols/brainmri-vit-model"
-model = ViTForImageClassification.from_pretrained(repository_id)
-feature_extractor = ViTImageProcessor.from_pretrained(repository_id)
+import base64
+from transformers import ViTForImageClassification, ViTImageProcessor
+from groq import Groq
+import numpy as np
+import torch.nn.functional as F
 
-# ------------------ ViT Inference Function ------------------
-def predict(image):
-    """
-    Given an image, perform inference using the ViT model to detect brain tumor.
-    Returns a human-readable diagnosis string.
-    """
-    # Convert to RGB and preprocess the image
-    image = image.convert("RGB")
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Set the device (GPU if available)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-    # Perform inference without gradient computation
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Get the predicted label and map to a diagnosis
-    logits = outputs.logits
-    predicted_label = logits.argmax(-1).item()
-    label_map = {0: "No", 1: "Yes"}
-    diagnosis = label_map[predicted_label]
-    
-    if diagnosis == "Yes":
-        return "The diagnosis indicates that you have a brain tumor."
-    else:
-        return "The diagnosis indicates that you do not have a brain tumor."
-
-# ------------------ Deepseek R1 Assistance Function ------------------
-def get_assistance_from_deepseek(diagnosis_text):
-    """
-    Given the diagnosis from the ViT model, call the Deepseek R1 model via the Groq API
-    to get additional recommendations and next steps.
-    """
+# ------------------ Logging ------------------
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-    api_key=os.getenv("API_KEY")
-    # Instantiate the Groq client with the provided API key
-    client = Groq(api_key=api_key)
-    
-    # Construct a prompt that includes the diagnosis and asks for detailed guidance
-    prompt = (
-        f"Based on the following diagnosis: '{diagnosis_text}', please provide next steps and "
-        "recommendations for the patient. Include whether to consult a specialist, if further tests "
-        "are needed, and any other immediate actions or lifestyle recommendations."
-    )
-    
-    messages = [
-        {
-            "role": "system",
-            "content": "You are a helpful medical assistant providing guidance after a brain tumor diagnosis."
-        },
-        {"role": "user", "content": prompt}
-    ]
-    
-    # Create the completion using the Deepseek R1 model (non-streaming for simplicity)
-    completion = client.chat.completions.create(
-        model="deepseek-r1-distill-llama-70b",
-        messages=messages,
-        temperature=0.6,
-        max_completion_tokens=4096,
-        top_p=0.95,
-        stream=False,
-        stop=None,
-    )
-    
-    # Extract the response text. (Depending on the API response format, adjust as needed.)
-    try:
-        assistance_text = completion.choices[0].message.content
-    except AttributeError:
-        # Fallback in case the structure is different
-        assistance_text = completion.choices[0].text
-    
-    return assistance_text
-
-# ------------------ Custom CSS for Styling ------------------
+# ------------------ Page config + CSS ------------------
+st.set_page_config(layout="wide", page_title="Brain MRI Tumor Detection")
 combined_css = """
     .main, .sidebar .sidebar-content { background-color: #1c1c1c; color: #f0f2f6; }
     .block-container { padding: 1rem 2rem; background-color: #333; border-radius: 10px; box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.5); }
@@ -119,13 +73,12 @@ combined_css = """
         margin-top: -20px;
         margin-bottom: 20px;
     }
+    .disclaimer { color: #ffcc66; font-weight: bold; text-align: center; margin-bottom: 12px; }
+    .small-muted { font-size:0.9rem; color:#cccccc; text-align:center; margin-top:8px; }
 """
-
-# ------------------ Streamlit App Configuration ------------------
-st.set_page_config(layout="wide")
 st.markdown(f"<style>{combined_css}</style>", unsafe_allow_html=True)
 
-# App Title and Description
+# ------------------ App header ------------------
 st.markdown(
     '<div class="title"><span class="colorful-text">Brain MRI</span> <span class="black-white-text">Tumor Detection</span></div>',
     unsafe_allow_html=True
@@ -135,37 +88,184 @@ st.markdown(
     unsafe_allow_html=True
 )
 
-# ------------------ Image Upload Section ------------------
+# Medical disclaimer (visible)
+st.markdown(
+    "<div class='disclaimer'>⚠️ This app is experimental and informational only. It is NOT a medical diagnosis. "
+    "If you have health concerns, consult a licensed medical professional. In emergencies call your local emergency number.</div>",
+    unsafe_allow_html=True
+)
+
+# ------------------ Model loading with graceful errors ------------------
+repository_id = "EnDevSols/brainmri-vit-model"
+
+model = None
+feature_extractor = None
+model_load_error = None
+
+try:
+    # Only attempt to load model if torch was imported successfully
+    if torch is None:
+        raise RuntimeError("Torch is not available in this environment.")
+    # Model loading can be slow; catch errors and show a friendly message later.
+    model = ViTForImageClassification.from_pretrained(repository_id)
+    feature_extractor = ViTImageProcessor.from_pretrained(repository_id)
+    logging.info("Model and feature extractor loaded successfully.")
+except Exception as e:
+    model_load_error = str(e)
+    logging.error("Failed to load model or feature extractor: %s\n%s", e, traceback.format_exc())
+
+# ------------------ Prediction function (returns label + confidence) ------------------
+def predict(image):
+    """
+    Given a PIL image, returns (diagnosis_label, confidence_float).
+    'diagnosis_label' is "Yes" (tumor) or "No" (no tumor).
+    'confidence_float' is between 0 and 1.
+    """
+    if model is None or feature_extractor is None:
+        raise RuntimeError("Model is not loaded.")
+    image = image.convert("RGB")
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    probs = F.softmax(logits, dim=-1).squeeze().cpu().numpy()
+    predicted_idx = int(np.argmax(probs))
+    confidence = float(probs[predicted_idx])
+    label_map = {0: "No", 1: "Yes"}
+    diagnosis = label_map.get(predicted_idx, "Unknown")
+    return diagnosis, confidence
+
+# ------------------ Deepseek (Groq) helper ------------------
+def get_assistance_from_deepseek(diagnosis_text):
+    """
+    Calls Groq Deepseek R1 with a safety-first prompt.
+    Returns a string. On error, returns a conservative fallback message.
+    """
+    api_key = os.getenv("API_KEY")
+    if not api_key:
+        logging.error("API_KEY environment variable not found for Groq client.")
+        return ("No assistance available because the Deepseek API key is not configured. "
+                "Please set the API_KEY environment variable.")
+    try:
+        client = Groq(api_key=api_key)
+    except Exception as e:
+        logging.error("Failed to instantiate Groq client: %s\n%s", e, traceback.format_exc())
+        return ("Assistance temporarily unavailable (failed to initialize model client). "
+                "Please try again later or contact support.")
+
+    # Safer prompt: require the assistant to include a clinician referral sentence
+    safety_sentence = "This information is informational only — seek evaluation from a licensed medical professional."
+    prompt = (
+        f"You are a cautious, safety-first medical assistant. Given the model-diagnosis below, "
+        "provide general, non-prescriptive information a patient could use to understand options. "
+        "Do NOT provide definitive medical diagnosis or treatment plans. ALWAYS include the sentence: "
+        f"'{safety_sentence}'\n\n"
+        f"Diagnosis text: '{diagnosis_text}'\n\n"
+        "Please list: (1) suggested questions a patient might ask a clinician, (2) common next diagnostic tests a clinician might consider (non-exhaustive), "
+        "and (3) immediate safety signs that should prompt emergency care. Keep the language simple and avoid prescriptive medical directives."
+    )
+
+    messages = [
+        {"role": "system", "content": "You are a careful medical assistant always advising a user to consult a clinician."},
+        {"role": "user", "content": prompt}
+    ]
+
+    try:
+        completion = client.chat.completions.create(
+            model="deepseek-r1-distill-llama-70b",
+            messages=messages,
+            temperature=0.6,
+            max_completion_tokens=1024,
+            top_p=0.95,
+            stream=False,
+            stop=None,
+        )
+        # Try different response shapes safely
+        assistance_text = ""
+        try:
+            assistance_text = completion.choices[0].message.content
+        except Exception:
+            try:
+                assistance_text = completion.choices[0].text
+            except Exception:
+                assistance_text = str(completion)
+
+        # Ensure the required safety sentence is present
+        if safety_sentence not in assistance_text:
+            assistance_text = safety_sentence + "\n\n" + assistance_text
+
+        return assistance_text
+    except Exception as e:
+        logging.error("Deepseek Groq call failed: %s\n%s", e, traceback.format_exc())
+        return ("Assistance is temporarily unavailable due to an error contacting the assistance model. "
+                "Please consult a licensed medical professional for evaluation. If you are experiencing severe or life-threatening symptoms, seek emergency care immediately.")
+
+# ------------------ Streamlit UI: image upload + inference flow ------------------
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    
-    # Resize image for display purposes
-    resized_image = image.resize((150, 150))
-    
-    # Convert image to base64 for HTML display
-    buffered = BytesIO()
-    resized_image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-
-    # Display the uploaded image in the center
-    st.markdown(
-        f"<div style='text-align: center;'><img src='data:image/jpeg;base64,{img_str}' alt='Uploaded Image' width='300'></div>",
-        unsafe_allow_html=True
-    )
-    
-    st.write("")
-    st.write("Processing the image...")
-
-    # ------------------ Step 1: Get Diagnosis from the ViT Model ------------------
-    diagnosis = predict(image)
-    st.markdown("### Diagnosis:")
-    st.write(diagnosis)
-    
-    # ------------------ Step 2: Get Further Assistance from Deepseek R1 ------------------
-    with st.spinner("Fetching additional guidance based on your diagnosis..."):
-        assistance = get_assistance_from_deepseek(diagnosis)
-    
-    st.markdown("### Next Steps and Recommendations:")
-    st.write(assistance)
+    try:
+        image = Image.open(uploaded_file)
+    except Exception as e:
+        st.error(f"Failed to open the uploaded file as an image: {e}")
+        image = None
+
+    if image is not None:
+        # Display resized thumbnail
+        resized_image = image.resize((150, 150))
+        buffered = BytesIO()
+        resized_image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        st.markdown(
+            f"<div style='text-align: center;'><img src='data:image/jpeg;base64,{img_str}' alt='Uploaded Image' width='300'></div>",
+            unsafe_allow_html=True
+        )
+
+        # Check model loaded
+        if model_load_error:
+            st.error("The model failed to load at startup. See logs for details.")
+            st.code(model_load_error)
+        else:
+            st.write("")
+            st.write("Processing the image...")
+
+            # Run prediction with try/except
+            try:
+                diagnosis, confidence = predict(image)
+                st.markdown("### Diagnosis (model prediction):")
+                st.write(f"**{diagnosis}** (confidence: **{confidence:.2%}**)")
+                st.markdown("_Model output is probabilistic and not a clinical diagnosis._", unsafe_allow_html=True)
+            except Exception as e:
+                st.error("Prediction failed: " + str(e))
+                logging.error("Prediction error: %s\n%s", e, traceback.format_exc())
+                diagnosis = None
+                confidence = None
+
+            # If we have a diagnosis, call Deepseek for additional guidance (with spinner)
+            if diagnosis is not None:
+                with st.spinner("Fetching additional guidance based on your diagnosis..."):
+                    assistance = get_assistance_from_deepseek(f"Diagnosis: {diagnosis} (confidence {confidence:.2%})")
+                st.markdown("### Next Steps and Recommendations:")
+                # Use st.write which keeps newlines and formatting reasonable.
+                st.write(assistance)
+
+# ------------------ If no file uploaded, show sample placeholder / instructions ------------------
+if uploaded_file is None:
+    st.markdown("<div class='small-muted'>Upload a brain MRI image (jpg/png) to get a model prediction and informational next steps. </div>", unsafe_allow_html=True)
+
+# ------------------ Helpful debug / info for developers (hidden by default) ------------------
+with st.expander("Developer info / Troubleshooting"):
+    st.markdown("**Model repository**: " + repository_id)
+    st.markdown("**Torch available**: " + ("Yes" if torch is not None else "No"))
+    st.markdown("**Model loaded**: " + ("Yes" if model is not None else "No"))
+    if model_load_error:
+        st.code(model_load_error)
+    st.markdown("**Environment**:")
+    st.write({
+        "CUDA available": torch.cuda.is_available() if torch is not None else False,
+        "API_KEY set for Groq": bool(os.getenv("API_KEY"))
+    })
+    st.markdown("**Notes:**\n- This app is for informational use only. Do not use it as a replacement for professional medical advice.")