Update cerebAI.py

cerebAI.py

@@ -11,34 +11,40 @@ import albumentations as A
 from albumentations.pytorch import ToTensorV2
 import os 
 import requests 
-import pydicom # REQUIRED FOR DICOM SUPPORT
+import pydicom 
 import io 
-import gc # For memory management
+import gc 
 
 # --- 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") # For Streamlit Cloud stability
+DEVICE = torch.device("cpu") 
 
-# --- MODEL LOADING ---
+# --- INITIALIZE SESSION STATE (CRITICAL FOR 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) ---
 @st.cache_resource
 def load_model(model_url, local_path):
     """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 as e:
-            st.error(f"FATAL ERROR: Could not download model. Check the URL. Error: {e}")
+        except Exception:
+            st.error(f"FATAL ERROR: Could not download model.")
             return None
 
     try:
@@ -48,11 +54,11 @@ def load_model(model_url, local_path):
         model.to(DEVICE)
         model.eval()
         return model
-    except Exception as e:
-        st.error(f"Failed to load model weights from cache. Error: {e}")
+    except Exception:
+        st.error(f"Failed to load model weights.")
         return None
 
-# --- HELPER FUNCTIONS ---
+# --- CORE LOGIC FUNCTIONS ---
 
 def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
     """Denormalizes a PyTorch tensor for matplotlib visualization."""
@@ -69,24 +75,19 @@ def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
 def preprocess_image(image_bytes: bytes, file_name: str) -> Tuple[Optional[torch.Tensor], Optional[np.ndarray]]:
     """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)
@@ -94,13 +95,10 @@ 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)
     
@@ -147,6 +145,46 @@ 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 --------------------
 # ==============================================================================
@@ -155,20 +193,20 @@ 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 (Sidebar or Main Area) ---
+    # --- INTERACTIVE CONTROLS ---
     st.markdown("### Analysis Controls")
     
     n_steps_slider = st.slider(
         'Integration Steps (Affects Accuracy & Speed)',
         min_value=5, 
         max_value=50, 
-        value=20, 
+        value=10, 
         step=5,
-        help="Higher steps (up to 50) provide a smoother, more accurate heatmap but use more CPU."
+        key="n_steps_slider_key", 
+        help="Lower steps (5-15) provide a smoother, more accurate heatmap but use more CPU."
     )
     st.markdown("---")
 
@@ -177,51 +215,46 @@ 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"]
+        type=["dcm", "dicom", "png", "jpg", "jpeg"],
+        key="file_uploader" 
     )
-
+    
+    # --- UI FLOW MANAGEMENT ---
     if uploaded_file is not None:
-        image_bytes = uploaded_file.read()
-        file_name = uploaded_file.name 
         
-        # 1. PROCESS IMAGE FIRST (Defines original_image_rgb)
-        input_tensor, original_image_rgb = preprocess_image(image_bytes, file_name) 
+        # 1. READ AND CACHE INPUT BYTES (Fastest Operation)
+        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)
+        # This part runs quickly and stabilizes the top half of the screen
+        _, original_image_rgb = preprocess_image(st.session_state.input_bytes, st.session_state.file_name) 
 
-        # --- DISPLAY AND RESULTS LAYOUT ---
+        # --- DISPLAY UPLOADED IMAGE AND BUTTON ---
         col1, col2 = st.columns(2) 
 
         with col1:
             st.subheader("Uploaded Image")
-            # 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)
-            
-            predicted_label = CLASS_LABELS[predicted_class_idx]
-            confidence_score = probabilities[predicted_class_idx]
+            st.image(original_image_rgb, use_container_width=True, caption=st.session_state.file_name) 
             
-            # Generate Attribution
-            heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
+            # --- THE GATED BUTTON (CRITICAL) ---
+            st.markdown("---")
+            if st.button("▶️ RUN FULL DIAGNOSIS & ATTRIBUTION", type="primary", use_container_width=True, key="run_button"):
+                # Call the heavy analysis function
+                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
             
-            # CRITICAL MEMORY MANAGEMENT
-            del input_tensor 
-            del output
-            gc.collect() 
-
             with col2:
                 st.subheader("Prediction Summary")
                 
                 st.metric(
                     label="Diagnosis", 
-                    value=predicted_label,
-                    delta=f"{confidence_score*100:.2f}% Confidence",
-                    delta_color='normal' 
+                    value=data['label'],
+                    delta=f"{data['confidence']*100:.2f}% Confidence",
+                    delta_color='normal'
                 )
                 
                 st.markdown("---")
@@ -229,7 +262,7 @@ if model is not None:
                 
                 prob_data = {
                     'Class': CLASS_LABELS,
-                    'Confidence': [f"{p:.4f}" for p in probabilities]
+                    'Confidence': [f"{p:.4f}" for p in data['probabilities']]
                 }
                 st.dataframe(prob_data, hide_index=True, use_container_width=True)
 
@@ -237,7 +270,7 @@ if model is not None:
             st.markdown("---")
             st.subheader("Model Interpretation (Integrated Gradients)")
             
-            fig = plot_heatmap_and_original(original_image_rgb, heatmap, predicted_label)
+            fig = plot_heatmap_and_original(data['image_rgb'], data['heatmap'], data['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.")