flicker fix rollback

cerebAI.py

@@ -15,31 +15,27 @@ import pydicom
 import io 
 import gc 
 
-# --- CONFIGURATION & CACHING ---
+# --- 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") 
 
-# --- INITIALIZE SESSION STATE (NEW) ---
-# This ensures variables persist between reruns and controls the UI flow.
-if 'output_ready' not in st.session_state:
-    st.session_state.output_ready = False
+# --- INITIALIZE SESSION STATE (NEW: for stable UI) ---
+if 'results_ready' not in st.session_state:
+    st.session_state.results_ready = False
     st.session_state.results = {}
-    st.session_state.original_image_rgb = None
-    st.session_state.input_tensor = None
+    st.session_state.file_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:
-            # ... (Download code) ...
             response = requests.get(model_url, stream=True)
             response.raise_for_status() 
             with open(local_path, "wb") as f:
@@ -58,60 +54,137 @@ 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
 
-# --- HELPER FUNCTIONS (UNCHANGED) ---
-# ... (denormalize_image, preprocess_image, generate_attribution, plot_heatmap_and_original functions here) ...
+# --- CORE HELPER FUNCTIONS (UNCHANGED) ---
 
-# -------------------- CORE COMPUTATION FUNCTION (NEW) --------------------
-
-def run_diagnosis(model, input_bytes, file_name, n_steps):
-    """
-    Handles all computation, run inside st.spinner.
-    Uses st.session_state to pass results out of the function.
-    """
-    try:
-        # Preprocess the uploaded image (The jittering part runs here, ONCE)
-        input_tensor, original_image_rgb = preprocess_image(input_bytes, file_name) 
+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:
+        tensor = tensor.detach() 
+    
+    mean, std = np.array([0.5, 0.5, 0.5]), np.array([0.5, 0.5, 0.5])
+    img = tensor.cpu().permute(1, 2, 0).numpy()
+    img = (img * std) + mean
+    return np.clip(img, 0, 1)
 
-        if input_tensor is None:
-            st.error("Could not process file. Ensure it is a valid DICOM/PNG/JPG.")
-            return
+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:
+            dcm = pydicom.dcmread(io.BytesIO(image_bytes))
+            pixel_array = dcm.pixel_array.astype(np.int16)
+            slope = dcm.RescaleSlope
+            intercept = dcm.RescaleIntercept
+            pixel_array = pixel_array * slope + intercept
             
-        # 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)
-        
-        # Generate Attribution (Heavy Calculation)
-        heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps)
-        
-        # --- CLEANUP (CRITICAL FOR RAM) ---
-        del input_tensor
-        del output
-        gc.collect() 
-        torch.cuda.empty_cache()
-
-        # --- STORE FINAL RESULTS ---
-        st.session_state.results = {
-            'label': CLASS_LABELS[predicted_class_idx],
-            'confidence': probabilities[predicted_class_idx],
-            'probabilities': probabilities,
-            'image_rgb': original_image_rgb,
-            'heatmap': heatmap,
-            'filename': file_name
-        }
-        st.session_state.output_ready = True
-
-    except Exception as e:
-        st.error(f"An unexpected error occurred during analysis: {e}")
-        st.session_state.output_ready = False
+            window_center = 40 
+            window_width = 150
+            min_hu = window_center - (window_width / 2)
+            max_hu = window_center + (window_width / 2)
+            
+            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)
+            
+        except Exception:
+            return None, None
+    else:
+        image_grayscale = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_GRAYSCALE)
+        if image_grayscale is None: return None, None
+            
+    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:
+    """Computes Integrated Gradients for the given input and class."""
+    target_class_int = int(predicted_class_idx) 
+    input_tensor.requires_grad_(True) 
+    
+    ig = IntegratedGradients(model)
+    baseline = torch.zeros_like(input_tensor).to(DEVICE)
+
+    attributions_ig = ig.attribute(
+        inputs=input_tensor,
+        baselines=baseline,
+        target=target_class_int,
+        n_steps=n_steps
+    )
+    
+    attributions_ig_vis = attributions_ig.squeeze(0).sum(dim=0).abs().cpu().detach().numpy()
+    
+    if attributions_ig_vis.max() > 0:
+        attributions_ig_vis = attributions_ig_vis / attributions_ig_vis.max()
+    
+    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) 
+
+    ax1.imshow(original_image_vis)
+    ax1.set_title("Original CT Scan", fontsize=14)
+    ax1.axis('off')
+
+    ax2.imshow(original_image_vis)
+    alpha_mask = heatmap * 0.7 + 0.3 
+
+    ax2.imshow(heatmap, cmap='jet', alpha=alpha_mask, vmin=0, vmax=1)
+    ax2.set_title(f"Interpretation: {predicted_label}", fontsize=14)
+    ax2.axis('off')
+    
+    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
+    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
+
+    # 2. 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)
+    
+    # 3. GENERATE ATTRIBUTION (HEAVY PART)
+    heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
+    
+    # 4. CRITICAL MEMORY MANAGEMENT
+    del input_tensor 
+    del output
+    gc.collect() 
+
+    # 5. SAVE FINAL RESULTS TO STATE
+    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() # Force a single, clean render of the results
 
 
 # ==============================================================================
-# -------------------- STREAMLIT FRONTEND (Final Clean Flow) --------------------
+# -------------------- STREAMLIT FRONTEND --------------------
 # ==============================================================================
 
 st.set_page_config(page_title="CerebAI: Stroke Prediction Dashboard", layout="wide")
@@ -124,62 +197,58 @@ if model is not None:
     # --- INPUT CONTROLS ---
     st.markdown("### Analysis Controls")
     
-    col_upload, col_steps = st.columns([2, 1])
-
-    with col_steps:
-        n_steps_slider = st.slider(
-            'Integration Steps (CPU Load)',
-            min_value=5, max_value=50, value=10, step=5,
-            help="Lower steps (5-15) for fast diagnosis; higher steps (30-50) for smoother visualization."
-        )
-    
-    with col_upload:
-        uploaded_file = st.file_uploader(
-            "Upload CT Scan Image", 
-            type=["dcm", "dicom", "png", "jpg", "jpeg"],
-            key="file_uploader" # Use key to manage state changes
-        )
-        
+    n_steps_slider = st.slider(
+        'Integration Steps (Affects Accuracy & Speed)',
+        min_value=5, 
+        max_value=50, 
+        value=10, 
+        step=5,
+        key="n_steps_slider_key", # Added key to prevent slider reset on rerun
+        help="Higher steps (up to 50) provide a smoother, more accurate heatmap but use more CPU."
+    )
     st.markdown("---")
+
+
+    # --- FILE UPLOAD ---
+    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" # Added key to manage file state
+    )
+    
     
-    # --- DISPLAY INPUT & RUN BUTTON ---
+    # --- UI FLOW MANAGEMENT ---
     if uploaded_file is not None:
-        
-        # NOTE: Read file bytes and display immediately (fast operation)
         image_bytes = uploaded_file.read()
         file_name = uploaded_file.name 
-        
-        # Display the file (use a generic path/bytes for fast display)
-        with col_upload:
-            st.subheader("Input Image")
-            st.image(image_bytes, use_container_width=True, caption=file_name)
-            
-            # --- THE GATING BUTTON ---
-            # This button calls the function that does the slow processing
-            if st.button("▶️ RUN FULL DIAGNOSIS", type="primary", use_container_width=True):
-                # We need to re-read the file and run the full logic inside the spinner
-                run_diagnosis(model, image_bytes, file_name, n_steps_slider)
-    
-    
-    # --- RESULTS DISPLAY ---
-    # This entire block only runs AFTER the button is clicked and the state is updated
-    if st.session_state.output_ready:
+
+        # Display the image immediately (fast operation, no lag)
+        st.subheader("Input Image")
+        st.image(image_bytes, use_container_width=True, caption=file_name) 
+
+        # --- GATED BUTTON ---
+        if st.button("▶️ RUN DIAGNOSIS & ATTRIBUTION", type="primary", use_container_width=True, key="run_button"):
+            with st.spinner(f'Calculating Integrated Gradients ({n_steps_slider} steps)...'):
+                # Execute the heavy logic and save to session state
+                perform_full_analysis(model, image_bytes, file_name, n_steps_slider)
+                
+    # --- RESULTS DISPLAY (ONLY RUNS AFTER BUTTON CLICK) ---
+    if st.session_state.results_ready:
         data = st.session_state.results
         
         st.markdown("---")
         st.markdown("### 2. Diagnosis and Interpretation")
 
-        # Stable layout for results
         col_metric, col_plot = st.columns([1, 2])
 
-        # --- METRICS ---
         with col_metric:
             st.subheader("Prediction Summary")
             st.metric(
-                label="Predicted Diagnosis", 
+                label="Diagnosis", 
                 value=data['label'],
                 delta=f"{data['confidence']*100:.2f}% Confidence",
-                delta_color='normal'
+                delta_color='normal' 
             )
             st.markdown("---")
             st.subheader("Confidence Breakdown")
@@ -189,10 +258,8 @@ if model is not None:
             }
             st.dataframe(prob_data, hide_index=True, use_container_width=True)
 
-        # --- PLOT ---
         with col_plot:
-            st.subheader("Visual Explanation (Integrated Gradients)")
-            
+            st.subheader("Visual Explanation")
             fig = plot_heatmap_and_original(data['image_rgb'], data['heatmap'], data['label'])
             st.pyplot(fig, clear_figure=True, use_container_width=True)