Update cerebAI.py

cerebAI.py

@@ -23,17 +23,17 @@ IMAGE_SIZE = 224
 DEVICE = torch.device("cpu") 
 
 # --- INITIALIZE SESSION STATE (CRITICAL FOR STABILITY) ---
-if 'results_ready' not in st.session_state:
-    st.session_state.results_ready = False
+if 'output_ready' not in st.session_state:
+    st.session_state.output_ready = False
     st.session_state.results = {}
-    st.session_state.input_bytes = None
     st.session_state.file_name = None
+    st.session_state.input_bytes = None
 
 
-# --- MODEL LOADING (UNCHANGED) ---
+# --- MODEL LOADING ---
 @st.cache_resource
 def load_model(model_url, local_path):
-    """Downloads model from URL if not cached, and loads the weights."""
+    # ... (Model loading logic remains the same) ...
     if not os.path.exists(local_path):
         st.info(f"Model not found locally. Downloading from remote repository...")
         try:
@@ -55,13 +55,12 @@ def load_model(model_url, local_path):
         model.eval()
         return model
     except Exception:
-        st.error(f"Failed to load model weights.")
+        st.error(f"Failed to load model weights from cache.")
         return None
 
-# --- CORE LOGIC FUNCTIONS ---
-
+# --- HELPER FUNCTIONS ---
+# ... (denormalize_image, preprocess_image, generate_attribution, plot_heatmap_and_original functions remain the same) ...
 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:
@@ -73,7 +72,6 @@ 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]]:
-    """Loads, processes, and normalizes image, handling DICOM or JPG/PNG."""
     
     if file_name.lower().endswith(('.dcm', '.dicom')):
         try:
@@ -105,7 +103,6 @@ def preprocess_image(image_bytes: bytes, file_name: str) -> Tuple[Optional[torch
     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:
-    """Computes Integrated Gradients for the given input and class."""
     target_class_int = int(predicted_class_idx) 
     input_tensor.requires_grad_(True) 
     
@@ -127,7 +124,6 @@ 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):
-    """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) 
 
@@ -145,7 +141,8 @@ 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) --------------------
+
+# -------------------- CORE DIAGNOSIS FUNCTION (GATED LOGIC) --------------------
 
 def perform_full_analysis(model, image_bytes, file_name, n_steps_slider):
     """Function called by the button to perform heavy computation and update state."""
@@ -196,7 +193,7 @@ st.markdown("---")
 model = load_model(HF_MODEL_URL, DOWNLOAD_MODEL_PATH)
 
 if model is not None:
-    # --- INTERACTIVE CONTROLS ---
+    # --- INPUT CONTROLS ---
     st.markdown("### Analysis Controls")
     
     n_steps_slider = st.slider(
@@ -206,7 +203,7 @@ if model is not None:
         value=10, 
         step=5,
         key="n_steps_slider_key", 
-        help="Lower steps (5-15) provide a smoother, more accurate heatmap but use more CPU."
+        help="Higher steps (up to 50) provide a smoother, more accurate heatmap but use more CPU."
     )
     st.markdown("---")
 
@@ -216,32 +213,33 @@ if model is not None:
     uploaded_file = st.file_uploader(
         "Choose a Dicom, PNG, JPG, or JPEG file", 
         type=["dcm", "dicom", "png", "jpg", "jpeg"],
-        key="file_uploader" 
+        key="file_uploader",
+        on_change=lambda: setattr(st.session_state, 'results_ready', False) # Reset results on new file upload
     )
-    
-    # --- UI FLOW MANAGEMENT ---
+
     if uploaded_file is not None:
         
         # 1. READ AND CACHE INPUT BYTES (Fastest Operation)
-        st.session_state.input_bytes = uploaded_file.read()
-        st.session_state.file_name = uploaded_file.name
+        image_bytes = uploaded_file.read()
+        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) 
+        # This part ensures the image appears instantly without waiting for the model
+        _, original_image_rgb = preprocess_image(image_bytes, file_name) 
 
         # --- DISPLAY UPLOADED IMAGE AND BUTTON ---
         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) 
             
-            # --- THE GATED BUTTON (CRITICAL) ---
+            # --- THE GATED BUTTON (CRITICAL FOR STABILITY) ---
             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)
+                # Execute the heavy logic and save to session state
+                perform_full_analysis(model, image_bytes, file_name, n_steps_slider)
         
         # --- RESULTS DISPLAY (STABLE) ---
         if st.session_state.results_ready:
@@ -273,4 +271,5 @@ if model is not None:
             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.")
+            st.success("Analysis Complete: The heatmap highlights the regions most critical to the diagnosis.")
+