Update cerebAI.py

cerebAI.py

@@ -11,40 +11,34 @@ import albumentations as A
 from albumentations.pytorch import ToTensorV2
 import os 
 import requests 
-import pydicom 
+import pydicom # REQUIRED FOR DICOM SUPPORT
 import io 
-import gc 
+import gc # For memory management
 
 # --- CONFIGURATION ---
 HF_MODEL_URL = "https://huggingface.co/arshenoy/cerebAI-stroke-model/resolve/main/best_model.pth" 
 DOWNLOAD_MODEL_PATH = "best_model_cache.pth"
 CLASS_LABELS = ['No Stroke', 'Ischemic Stroke', 'Hemorrhagic Stroke']
 IMAGE_SIZE = 224
-DEVICE = torch.device("cpu") 
+DEVICE = torch.device("cpu") # For Streamlit Cloud stability
 
-# --- INITIALIZE SESSION STATE (CRITICAL FOR UI STABILITY) ---
-if 'results_ready' not in st.session_state:
-    st.session_state.results_ready = False
-    st.session_state.results = {}
-    st.session_state.input_bytes = None
-    st.session_state.file_name = None
-
-
-# --- MODEL LOADING (UNCHANGED) ---
+# --- MODEL LOADING ---
 @st.cache_resource
 def load_model(model_url, local_path):
-    # ... (Model download and loading logic) ...
+    """Downloads model from URL if not cached, and loads the weights."""
+    
     if not os.path.exists(local_path):
         st.info(f"Model not found locally. Downloading from remote repository...")
         try:
             response = requests.get(model_url, stream=True)
             response.raise_for_status() 
+            
             with open(local_path, "wb") as f:
                 for chunk in response.iter_content(chunk_size=8192):
                     f.write(chunk)
             st.success("Model download complete!")
-        except Exception:
-            st.error(f"FATAL ERROR: Could not download model.")
+        except Exception as e:
+            st.error(f"FATAL ERROR: Could not download model. Check the URL. Error: {e}")
             return None
 
     try:
@@ -54,12 +48,14 @@ def load_model(model_url, local_path):
         model.to(DEVICE)
         model.eval()
         return model
-    except Exception:
-        st.error(f"Failed to load model weights.")
+    except Exception as e:
+        st.error(f"Failed to load model weights from cache. Error: {e}")
         return None
 
-# --- CORE LOGIC FUNCTIONS (UNCHANGED) ---
+# --- HELPER FUNCTIONS ---
+
 def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
+    """Denormalizes a PyTorch tensor for matplotlib visualization."""
     if tensor.ndim == 4:
         tensor = tensor.squeeze(0).detach() 
     else:
@@ -71,20 +67,26 @@ def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
     return np.clip(img, 0, 1)
 
 def preprocess_image(image_bytes: bytes, file_name: str) -> Tuple[Optional[torch.Tensor], Optional[np.ndarray]]:
-    # ... (DICOM and preprocessing logic remains the same) ...
+    """Loads, processes, and normalizes image, handling DICOM or JPG/PNG."""
+    
+    # 1. READ IMAGE DATA (Handles DICOM vs Standard formats)
     if file_name.lower().endswith(('.dcm', '.dicom')):
         try:
             dcm = pydicom.dcmread(io.BytesIO(image_bytes))
+            
+            # FIX: Convert to Hounsfield Units (HU)
             pixel_array = dcm.pixel_array.astype(np.int16)
             slope = dcm.RescaleSlope
             intercept = dcm.RescaleIntercept
             pixel_array = pixel_array * slope + intercept
             
+            # Apply Standard Brain Window (-100 HU to 150 HU)
             window_center = 40 
             window_width = 150
             min_hu = window_center - (window_width / 2)
             max_hu = window_center + (window_width / 2)
             
+            # Apply the windowing transformation and scale to 0-255
             pixel_array[pixel_array < min_hu] = min_hu
             pixel_array[pixel_array > max_hu] = max_hu
             image_grayscale = ((pixel_array - min_hu) / (max_hu - min_hu) * 255).astype(np.uint8)
@@ -92,17 +94,20 @@ def preprocess_image(image_bytes: bytes, file_name: str) -> Tuple[Optional[torch
         except Exception:
             return None, None
     else:
+        # Read standard image (PNG/JPG)
         image_grayscale = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_GRAYSCALE)
         if image_grayscale is None: return None, None
             
+    # 2. STANDARD PREPROCESSING
     image_rgb = cv2.cvtColor(cv2.resize(image_grayscale, (IMAGE_SIZE, IMAGE_SIZE)), cv2.COLOR_GRAY2RGB)
+    
     image_norm = (image_rgb.astype(np.float32) / 255.0 - 0.5) / 0.5
     input_tensor = torch.tensor(image_norm, dtype=torch.float).permute(2, 0, 1).unsqueeze(0)
     
     return input_tensor.to(DEVICE), image_rgb
 
 def generate_attribution(model: nn.Module, input_tensor: torch.Tensor, predicted_class_idx: int, n_steps: int = 20) -> np.ndarray:
-    # ... (Attribution logic remains the same) ...
+    """Computes Integrated Gradients for the given input and class."""
     target_class_int = int(predicted_class_idx) 
     input_tensor.requires_grad_(True) 
     
@@ -124,7 +129,7 @@ def generate_attribution(model: nn.Module, input_tensor: torch.Tensor, predicted
     return attributions_ig_vis
 
 def plot_heatmap_and_original(original_image: np.ndarray, heatmap: np.ndarray, predicted_label: str):
-    # ... (Plotting logic remains the same) ...
+    """Creates a Matplotlib figure for visualization."""
     fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) 
     original_image_vis = (original_image.astype(np.float32) / 255.0) 
 
@@ -142,46 +147,6 @@ def plot_heatmap_and_original(original_image: np.ndarray, heatmap: np.ndarray, p
     plt.tight_layout()
     return fig
 
-# -------------------- CORE DIAGNOSIS FUNCTION (TO BE CALLED BY BUTTON) --------------------
-
-def perform_full_analysis(model, image_bytes, file_name, n_steps_slider):
-    """Function called by the button to perform heavy computation and update state."""
-    
-    # 1. PREPARE INPUTS (The Jitter-Causing step)
-    input_tensor, original_image_rgb = preprocess_image(image_bytes, file_name) 
-
-    if input_tensor is None:
-        st.error("Could not process file. Ensure it is a valid DICOM/PNG/JPG.")
-        st.session_state.results_ready = False
-        return
-    
-    with st.spinner(f'Calculating Integrated Gradients ({n_steps_slider} steps)...'):
-
-        # 2. PREDICT & GENERATE ATTRIBUTION
-        with torch.no_grad():
-            output = model(input_tensor)
-            probabilities = torch.softmax(output, dim=1).squeeze(0).cpu().numpy()
-            predicted_class_idx = np.argmax(probabilities)
-        
-        heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
-        
-        # 3. CRITICAL MEMORY MANAGEMENT
-        del input_tensor 
-        del output
-        gc.collect() 
-
-    # 4. STORE FINAL RESULTS TO STATE AND RERUN
-    st.session_state.results = {
-        'label': CLASS_LABELS[predicted_class_idx],
-        'confidence': probabilities[predicted_class_idx],
-        'probabilities': probabilities,
-        'image_rgb': original_image_rgb,
-        'heatmap': heatmap
-    }
-    st.session_state.results_ready = True
-    st.rerun() 
-
-
 # ==============================================================================
 # -------------------- STREAMLIT FRONTEND --------------------
 # ==============================================================================
@@ -190,19 +155,19 @@ st.set_page_config(page_title="CerebAI: Stroke Prediction Dashboard", layout="wi
 st.title("CerebAI: AI-Powered Stroke Detection")
 st.markdown("---")
 
+# FIX: Load the model using the download mechanism
 model = load_model(HF_MODEL_URL, DOWNLOAD_MODEL_PATH)
 
 if model is not None:
-    # --- INTERACTIVE CONTROLS ---
+    # --- INTERACTIVE CONTROLS (Sidebar or Main Area) ---
     st.markdown("### Analysis Controls")
     
     n_steps_slider = st.slider(
         'Integration Steps (Affects Accuracy & Speed)',
         min_value=5, 
         max_value=50, 
-        value=10, 
+        value=20, 
         step=5,
-        key="n_steps_slider_key", 
         help="Higher steps (up to 50) provide a smoother, more accurate heatmap but use more CPU."
     )
     st.markdown("---")
@@ -212,43 +177,50 @@ if model is not None:
     st.markdown("### Upload CT Scan Image")
     uploaded_file = st.file_uploader(
         "Choose a Dicom, PNG, JPG, or JPEG file", 
-        type=["dcm", "dicom", "png", "jpg", "jpeg"],
-        key="file_uploader" 
+        type=["dcm", "dicom", "png", "jpg", "jpeg"]
     )
-    
-    # --- UI FLOW MANAGEMENT ---
+
     if uploaded_file is not None:
+        image_bytes = uploaded_file.read()
+        file_name = uploaded_file.name 
         
-        # Cache file bytes and name to session state to prevent corruption on rerun
-        st.session_state.input_bytes = uploaded_file.read()
-        st.session_state.file_name = uploaded_file.name
-        
-        # Preprocess image for DISPLAY ONLY (fast operation, no lag)
-        _, original_image_rgb = preprocess_image(st.session_state.input_bytes, st.session_state.file_name) 
+        # 1. PROCESS IMAGE FIRST (Defines original_image_rgb)
+        input_tensor, original_image_rgb = preprocess_image(image_bytes, file_name) 
 
-        # --- DISPLAY UPLOADED IMAGE AND BUTTON ---
+        # --- DISPLAY AND RESULTS LAYOUT ---
         col1, col2 = st.columns(2) 
 
         with col1:
             st.subheader("Uploaded Image")
-            st.image(original_image_rgb, use_container_width=True, caption=st.session_state.file_name) 
+            # Display the processed NumPy array
+            st.image(original_image_rgb, use_container_width=True, caption=file_name) 
+
+        # Run Prediction and Attribution
+        if input_tensor is not None:
+            # Predict
+            with torch.no_grad():
+                output = model(input_tensor)
+                probabilities = torch.softmax(output, dim=1).squeeze(0).cpu().numpy()
+                predicted_class_idx = np.argmax(probabilities)
             
-            # --- THE GATED BUTTON ---
-            # This button will trigger the heavy computation
-            if st.button("▶️ RUN FULL DIAGNOSIS & ATTRIBUTION", type="primary", use_container_width=True, key="run_button"):
-                perform_full_analysis(model, st.session_state.input_bytes, st.session_state.file_name, n_steps_slider)
-        
-        # --- RESULTS DISPLAY (STABLE) ---
-        if st.session_state.results_ready:
-            data = st.session_state.results
+            predicted_label = CLASS_LABELS[predicted_class_idx]
+            confidence_score = probabilities[predicted_class_idx]
+            
+            # Generate Attribution
+            heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
             
+            # CRITICAL MEMORY MANAGEMENT
+            del input_tensor 
+            del output
+            gc.collect() 
+
             with col2:
                 st.subheader("Prediction Summary")
                 
                 st.metric(
                     label="Diagnosis", 
-                    value=data['label'],
-                    delta=f"{data['confidence']*100:.2f}% Confidence",
+                    value=predicted_label,
+                    delta=f"{confidence_score*100:.2f}% Confidence",
                     delta_color='normal' 
                 )
                 
@@ -257,7 +229,7 @@ if model is not None:
                 
                 prob_data = {
                     'Class': CLASS_LABELS,
-                    'Confidence': [f"{p:.4f}" for p in data['probabilities']]
+                    'Confidence': [f"{p:.4f}" for p in probabilities]
                 }
                 st.dataframe(prob_data, hide_index=True, use_container_width=True)
 
@@ -265,7 +237,7 @@ if model is not None:
             st.markdown("---")
             st.subheader("Model Interpretation (Integrated Gradients)")
             
-            fig = plot_heatmap_and_original(data['image_rgb'], data['heatmap'], data['label'])
+            fig = plot_heatmap_and_original(original_image_rgb, heatmap, predicted_label)
             st.pyplot(fig, clear_figure=True, use_container_width=True) 
 
             st.success("Analysis Complete: The heatmap highlights the regions most critical to the diagnosis.")